mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-02-19 07:44:47 +08:00
509 lines
23 KiB
Python
509 lines
23 KiB
Python
import os
|
|
import glob
|
|
import logging
|
|
import importlib
|
|
from tqdm import tqdm
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from core.prefetch_dataloader import PrefetchDataLoader, CPUPrefetcher
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
import torchvision
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
from core.lr_scheduler import MultiStepRestartLR, CosineAnnealingRestartLR
|
|
from core.loss import AdversarialLoss, PerceptualLoss, LPIPSLoss
|
|
from core.dataset import TrainDataset
|
|
|
|
from model.modules.flow_comp_raft import RAFT_bi, FlowLoss, EdgeLoss
|
|
from model.recurrent_flow_completion import RecurrentFlowCompleteNet
|
|
|
|
from RAFT.utils.flow_viz_pt import flow_to_image
|
|
|
|
|
|
class Trainer:
|
|
def __init__(self, config):
|
|
self.config = config
|
|
self.epoch = 0
|
|
self.iteration = 0
|
|
self.num_local_frames = config['train_data_loader']['num_local_frames']
|
|
self.num_ref_frames = config['train_data_loader']['num_ref_frames']
|
|
|
|
# setup data set and data loader
|
|
self.train_dataset = TrainDataset(config['train_data_loader'])
|
|
|
|
self.train_sampler = None
|
|
self.train_args = config['trainer']
|
|
if config['distributed']:
|
|
self.train_sampler = DistributedSampler(
|
|
self.train_dataset,
|
|
num_replicas=config['world_size'],
|
|
rank=config['global_rank'])
|
|
|
|
dataloader_args = dict(
|
|
dataset=self.train_dataset,
|
|
batch_size=self.train_args['batch_size'] // config['world_size'],
|
|
shuffle=(self.train_sampler is None),
|
|
num_workers=self.train_args['num_workers'],
|
|
sampler=self.train_sampler,
|
|
drop_last=True)
|
|
|
|
self.train_loader = PrefetchDataLoader(self.train_args['num_prefetch_queue'], **dataloader_args)
|
|
self.prefetcher = CPUPrefetcher(self.train_loader)
|
|
|
|
# set loss functions
|
|
self.adversarial_loss = AdversarialLoss(type=self.config['losses']['GAN_LOSS'])
|
|
self.adversarial_loss = self.adversarial_loss.to(self.config['device'])
|
|
self.l1_loss = nn.L1Loss()
|
|
# self.perc_loss = PerceptualLoss(
|
|
# layer_weights={'conv3_4': 0.25, 'conv4_4': 0.25, 'conv5_4': 0.5},
|
|
# use_input_norm=True,
|
|
# range_norm=True,
|
|
# criterion='l1'
|
|
# ).to(self.config['device'])
|
|
|
|
if self.config['losses']['perceptual_weight'] > 0:
|
|
self.perc_loss = LPIPSLoss(use_input_norm=True, range_norm=True).to(self.config['device'])
|
|
|
|
# self.flow_comp_loss = FlowCompletionLoss().to(self.config['device'])
|
|
# self.flow_comp_loss = FlowCompletionLoss(self.config['device'])
|
|
|
|
# set raft
|
|
self.fix_raft = RAFT_bi(device = self.config['device'])
|
|
self.fix_flow_complete = RecurrentFlowCompleteNet('/mnt/lustre/sczhou/VQGANs/CodeMOVI/experiments_model/recurrent_flow_completion_v5_train_flowcomp_v5/gen_760000.pth')
|
|
for p in self.fix_flow_complete.parameters():
|
|
p.requires_grad = False
|
|
self.fix_flow_complete.to(self.config['device'])
|
|
self.fix_flow_complete.eval()
|
|
|
|
# self.flow_loss = FlowLoss()
|
|
|
|
# setup models including generator and discriminator
|
|
net = importlib.import_module('model.' + config['model']['net'])
|
|
self.netG = net.InpaintGenerator()
|
|
# print(self.netG)
|
|
self.netG = self.netG.to(self.config['device'])
|
|
if not self.config['model'].get('no_dis', False):
|
|
if self.config['model'].get('dis_2d', False):
|
|
self.netD = net.Discriminator_2D(
|
|
in_channels=3,
|
|
use_sigmoid=config['losses']['GAN_LOSS'] != 'hinge')
|
|
else:
|
|
self.netD = net.Discriminator(
|
|
in_channels=3,
|
|
use_sigmoid=config['losses']['GAN_LOSS'] != 'hinge')
|
|
self.netD = self.netD.to(self.config['device'])
|
|
|
|
self.interp_mode = self.config['model']['interp_mode']
|
|
# setup optimizers and schedulers
|
|
self.setup_optimizers()
|
|
self.setup_schedulers()
|
|
self.load()
|
|
|
|
if config['distributed']:
|
|
self.netG = DDP(self.netG,
|
|
device_ids=[self.config['local_rank']],
|
|
output_device=self.config['local_rank'],
|
|
broadcast_buffers=True,
|
|
find_unused_parameters=True)
|
|
if not self.config['model']['no_dis']:
|
|
self.netD = DDP(self.netD,
|
|
device_ids=[self.config['local_rank']],
|
|
output_device=self.config['local_rank'],
|
|
broadcast_buffers=True,
|
|
find_unused_parameters=False)
|
|
|
|
# set summary writer
|
|
self.dis_writer = None
|
|
self.gen_writer = None
|
|
self.summary = {}
|
|
if self.config['global_rank'] == 0 or (not config['distributed']):
|
|
if not self.config['model']['no_dis']:
|
|
self.dis_writer = SummaryWriter(
|
|
os.path.join(config['save_dir'], 'dis'))
|
|
self.gen_writer = SummaryWriter(
|
|
os.path.join(config['save_dir'], 'gen'))
|
|
|
|
def setup_optimizers(self):
|
|
"""Set up optimizers."""
|
|
backbone_params = []
|
|
for name, param in self.netG.named_parameters():
|
|
if param.requires_grad:
|
|
backbone_params.append(param)
|
|
else:
|
|
print(f'Params {name} will not be optimized.')
|
|
|
|
optim_params = [
|
|
{
|
|
'params': backbone_params,
|
|
'lr': self.config['trainer']['lr']
|
|
},
|
|
]
|
|
|
|
self.optimG = torch.optim.Adam(optim_params,
|
|
betas=(self.config['trainer']['beta1'],
|
|
self.config['trainer']['beta2']))
|
|
|
|
if not self.config['model']['no_dis']:
|
|
self.optimD = torch.optim.Adam(
|
|
self.netD.parameters(),
|
|
lr=self.config['trainer']['lr'],
|
|
betas=(self.config['trainer']['beta1'],
|
|
self.config['trainer']['beta2']))
|
|
|
|
def setup_schedulers(self):
|
|
"""Set up schedulers."""
|
|
scheduler_opt = self.config['trainer']['scheduler']
|
|
scheduler_type = scheduler_opt.pop('type')
|
|
|
|
if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']:
|
|
self.scheG = MultiStepRestartLR(
|
|
self.optimG,
|
|
milestones=scheduler_opt['milestones'],
|
|
gamma=scheduler_opt['gamma'])
|
|
if not self.config['model']['no_dis']:
|
|
self.scheD = MultiStepRestartLR(
|
|
self.optimD,
|
|
milestones=scheduler_opt['milestones'],
|
|
gamma=scheduler_opt['gamma'])
|
|
elif scheduler_type == 'CosineAnnealingRestartLR':
|
|
self.scheG = CosineAnnealingRestartLR(
|
|
self.optimG,
|
|
periods=scheduler_opt['periods'],
|
|
restart_weights=scheduler_opt['restart_weights'],
|
|
eta_min=scheduler_opt['eta_min'])
|
|
if not self.config['model']['no_dis']:
|
|
self.scheD = CosineAnnealingRestartLR(
|
|
self.optimD,
|
|
periods=scheduler_opt['periods'],
|
|
restart_weights=scheduler_opt['restart_weights'],
|
|
eta_min=scheduler_opt['eta_min'])
|
|
else:
|
|
raise NotImplementedError(
|
|
f'Scheduler {scheduler_type} is not implemented yet.')
|
|
|
|
def update_learning_rate(self):
|
|
"""Update learning rate."""
|
|
self.scheG.step()
|
|
if not self.config['model']['no_dis']:
|
|
self.scheD.step()
|
|
|
|
def get_lr(self):
|
|
"""Get current learning rate."""
|
|
return self.optimG.param_groups[0]['lr']
|
|
|
|
def add_summary(self, writer, name, val):
|
|
"""Add tensorboard summary."""
|
|
if name not in self.summary:
|
|
self.summary[name] = 0
|
|
self.summary[name] += val
|
|
n = self.train_args['log_freq']
|
|
if writer is not None and self.iteration % n == 0:
|
|
writer.add_scalar(name, self.summary[name] / n, self.iteration)
|
|
self.summary[name] = 0
|
|
|
|
def load(self):
|
|
"""Load netG (and netD)."""
|
|
# get the latest checkpoint
|
|
model_path = self.config['save_dir']
|
|
# TODO: add resume name
|
|
if os.path.isfile(os.path.join(model_path, 'latest.ckpt')):
|
|
latest_epoch = open(os.path.join(model_path, 'latest.ckpt'),
|
|
'r').read().splitlines()[-1]
|
|
else:
|
|
ckpts = [
|
|
os.path.basename(i).split('.pth')[0]
|
|
for i in glob.glob(os.path.join(model_path, '*.pth'))
|
|
]
|
|
ckpts.sort()
|
|
latest_epoch = ckpts[-1][4:] if len(ckpts) > 0 else None
|
|
|
|
if latest_epoch is not None:
|
|
gen_path = os.path.join(model_path,
|
|
f'gen_{int(latest_epoch):06d}.pth')
|
|
dis_path = os.path.join(model_path,
|
|
f'dis_{int(latest_epoch):06d}.pth')
|
|
opt_path = os.path.join(model_path,
|
|
f'opt_{int(latest_epoch):06d}.pth')
|
|
|
|
if self.config['global_rank'] == 0:
|
|
print(f'Loading model from {gen_path}...')
|
|
dataG = torch.load(gen_path, map_location=self.config['device'])
|
|
self.netG.load_state_dict(dataG)
|
|
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
|
dataD = torch.load(dis_path, map_location=self.config['device'])
|
|
self.netD.load_state_dict(dataD)
|
|
|
|
data_opt = torch.load(opt_path, map_location=self.config['device'])
|
|
self.optimG.load_state_dict(data_opt['optimG'])
|
|
# self.scheG.load_state_dict(data_opt['scheG'])
|
|
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
|
self.optimD.load_state_dict(data_opt['optimD'])
|
|
# self.scheD.load_state_dict(data_opt['scheD'])
|
|
self.epoch = data_opt['epoch']
|
|
self.iteration = data_opt['iteration']
|
|
else:
|
|
gen_path = self.config['trainer'].get('gen_path', None)
|
|
dis_path = self.config['trainer'].get('dis_path', None)
|
|
opt_path = self.config['trainer'].get('opt_path', None)
|
|
if gen_path is not None:
|
|
if self.config['global_rank'] == 0:
|
|
print(f'Loading Gen-Net from {gen_path}...')
|
|
dataG = torch.load(gen_path, map_location=self.config['device'])
|
|
self.netG.load_state_dict(dataG)
|
|
|
|
if dis_path is not None and not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
|
if self.config['global_rank'] == 0:
|
|
print(f'Loading Dis-Net from {dis_path}...')
|
|
dataD = torch.load(dis_path, map_location=self.config['device'])
|
|
self.netD.load_state_dict(dataD)
|
|
if opt_path is not None:
|
|
data_opt = torch.load(opt_path, map_location=self.config['device'])
|
|
self.optimG.load_state_dict(data_opt['optimG'])
|
|
self.scheG.load_state_dict(data_opt['scheG'])
|
|
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
|
self.optimD.load_state_dict(data_opt['optimD'])
|
|
self.scheD.load_state_dict(data_opt['scheD'])
|
|
else:
|
|
if self.config['global_rank'] == 0:
|
|
print('Warnning: There is no trained model found.'
|
|
'An initialized model will be used.')
|
|
|
|
def save(self, it):
|
|
"""Save parameters every eval_epoch"""
|
|
if self.config['global_rank'] == 0:
|
|
# configure path
|
|
gen_path = os.path.join(self.config['save_dir'],
|
|
f'gen_{it:06d}.pth')
|
|
dis_path = os.path.join(self.config['save_dir'],
|
|
f'dis_{it:06d}.pth')
|
|
opt_path = os.path.join(self.config['save_dir'],
|
|
f'opt_{it:06d}.pth')
|
|
print(f'\nsaving model to {gen_path} ...')
|
|
|
|
# remove .module for saving
|
|
if isinstance(self.netG, torch.nn.DataParallel) or isinstance(self.netG, DDP):
|
|
netG = self.netG.module
|
|
if not self.config['model']['no_dis']:
|
|
netD = self.netD.module
|
|
else:
|
|
netG = self.netG
|
|
if not self.config['model']['no_dis']:
|
|
netD = self.netD
|
|
|
|
# save checkpoints
|
|
torch.save(netG.state_dict(), gen_path)
|
|
if not self.config['model']['no_dis']:
|
|
torch.save(netD.state_dict(), dis_path)
|
|
torch.save(
|
|
{
|
|
'epoch': self.epoch,
|
|
'iteration': self.iteration,
|
|
'optimG': self.optimG.state_dict(),
|
|
'optimD': self.optimD.state_dict(),
|
|
'scheG': self.scheG.state_dict(),
|
|
'scheD': self.scheD.state_dict()
|
|
}, opt_path)
|
|
else:
|
|
torch.save(
|
|
{
|
|
'epoch': self.epoch,
|
|
'iteration': self.iteration,
|
|
'optimG': self.optimG.state_dict(),
|
|
'scheG': self.scheG.state_dict()
|
|
}, opt_path)
|
|
|
|
latest_path = os.path.join(self.config['save_dir'], 'latest.ckpt')
|
|
os.system(f"echo {it:06d} > {latest_path}")
|
|
|
|
def train(self):
|
|
"""training entry"""
|
|
pbar = range(int(self.train_args['iterations']))
|
|
if self.config['global_rank'] == 0:
|
|
pbar = tqdm(pbar,
|
|
initial=self.iteration,
|
|
dynamic_ncols=True,
|
|
smoothing=0.01)
|
|
|
|
os.makedirs('logs', exist_ok=True)
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(filename)s[line:%(lineno)d]"
|
|
"%(levelname)s %(message)s",
|
|
datefmt="%a, %d %b %Y %H:%M:%S",
|
|
filename=f"logs/{self.config['save_dir'].split('/')[-1]}.log",
|
|
filemode='w')
|
|
|
|
while True:
|
|
self.epoch += 1
|
|
self.prefetcher.reset()
|
|
if self.config['distributed']:
|
|
self.train_sampler.set_epoch(self.epoch)
|
|
self._train_epoch(pbar)
|
|
if self.iteration > self.train_args['iterations']:
|
|
break
|
|
print('\nEnd training....')
|
|
|
|
def _train_epoch(self, pbar):
|
|
"""Process input and calculate loss every training epoch"""
|
|
device = self.config['device']
|
|
train_data = self.prefetcher.next()
|
|
while train_data is not None:
|
|
self.iteration += 1
|
|
frames, masks, flows_f, flows_b, _ = train_data
|
|
frames, masks = frames.to(device), masks.to(device).float()
|
|
l_t = self.num_local_frames
|
|
b, t, c, h, w = frames.size()
|
|
gt_local_frames = frames[:, :l_t, ...]
|
|
local_masks = masks[:, :l_t, ...].contiguous()
|
|
|
|
masked_frames = frames * (1 - masks)
|
|
masked_local_frames = masked_frames[:, :l_t, ...]
|
|
# get gt optical flow
|
|
if flows_f[0] == 'None' or flows_b[0] == 'None':
|
|
gt_flows_bi = self.fix_raft(gt_local_frames)
|
|
else:
|
|
gt_flows_bi = (flows_f.to(device), flows_b.to(device))
|
|
|
|
# ---- complete flow ----
|
|
pred_flows_bi, _ = self.fix_flow_complete.forward_bidirect_flow(gt_flows_bi, local_masks)
|
|
pred_flows_bi = self.fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, local_masks)
|
|
# pred_flows_bi = gt_flows_bi
|
|
|
|
# ---- image propagation ----
|
|
prop_imgs, updated_local_masks = self.netG.module.img_propagation(masked_local_frames, pred_flows_bi, local_masks, interpolation=self.interp_mode)
|
|
updated_masks = masks.clone()
|
|
updated_masks[:, :l_t, ...] = updated_local_masks.view(b, l_t, 1, h, w)
|
|
updated_frames = masked_frames.clone()
|
|
prop_local_frames = gt_local_frames * (1-local_masks) + prop_imgs.view(b, l_t, 3, h, w) * local_masks # merge
|
|
updated_frames[:, :l_t, ...] = prop_local_frames
|
|
|
|
# ---- feature propagation + Transformer ----
|
|
pred_imgs = self.netG(updated_frames, pred_flows_bi, masks, updated_masks, l_t)
|
|
pred_imgs = pred_imgs.view(b, -1, c, h, w)
|
|
|
|
# get the local frames
|
|
pred_local_frames = pred_imgs[:, :l_t, ...]
|
|
comp_local_frames = gt_local_frames * (1. - local_masks) + pred_local_frames * local_masks
|
|
comp_imgs = frames * (1. - masks) + pred_imgs * masks
|
|
|
|
gen_loss = 0
|
|
dis_loss = 0
|
|
# optimize net_g
|
|
if not self.config['model']['no_dis']:
|
|
for p in self.netD.parameters():
|
|
p.requires_grad = False
|
|
|
|
self.optimG.zero_grad()
|
|
|
|
# generator l1 loss
|
|
hole_loss = self.l1_loss(pred_imgs * masks, frames * masks)
|
|
hole_loss = hole_loss / torch.mean(masks) * self.config['losses']['hole_weight']
|
|
gen_loss += hole_loss
|
|
self.add_summary(self.gen_writer, 'loss/hole_loss', hole_loss.item())
|
|
|
|
valid_loss = self.l1_loss(pred_imgs * (1 - masks), frames * (1 - masks))
|
|
valid_loss = valid_loss / torch.mean(1-masks) * self.config['losses']['valid_weight']
|
|
gen_loss += valid_loss
|
|
self.add_summary(self.gen_writer, 'loss/valid_loss', valid_loss.item())
|
|
|
|
# perceptual loss
|
|
if self.config['losses']['perceptual_weight'] > 0:
|
|
perc_loss = self.perc_loss(pred_imgs.view(-1,3,h,w), frames.view(-1,3,h,w))[0] * self.config['losses']['perceptual_weight']
|
|
gen_loss += perc_loss
|
|
self.add_summary(self.gen_writer, 'loss/perc_loss', perc_loss.item())
|
|
|
|
# gan loss
|
|
if not self.config['model']['no_dis']:
|
|
# generator adversarial loss
|
|
gen_clip = self.netD(comp_imgs)
|
|
gan_loss = self.adversarial_loss(gen_clip, True, False)
|
|
gan_loss = gan_loss * self.config['losses']['adversarial_weight']
|
|
gen_loss += gan_loss
|
|
self.add_summary(self.gen_writer, 'loss/gan_loss', gan_loss.item())
|
|
gen_loss.backward()
|
|
self.optimG.step()
|
|
|
|
if not self.config['model']['no_dis']:
|
|
# optimize net_d
|
|
for p in self.netD.parameters():
|
|
p.requires_grad = True
|
|
self.optimD.zero_grad()
|
|
|
|
# discriminator adversarial loss
|
|
real_clip = self.netD(frames)
|
|
fake_clip = self.netD(comp_imgs.detach())
|
|
dis_real_loss = self.adversarial_loss(real_clip, True, True)
|
|
dis_fake_loss = self.adversarial_loss(fake_clip, False, True)
|
|
dis_loss += (dis_real_loss + dis_fake_loss) / 2
|
|
self.add_summary(self.dis_writer, 'loss/dis_vid_real', dis_real_loss.item())
|
|
self.add_summary(self.dis_writer, 'loss/dis_vid_fake', dis_fake_loss.item())
|
|
dis_loss.backward()
|
|
self.optimD.step()
|
|
|
|
self.update_learning_rate()
|
|
|
|
# write image to tensorboard
|
|
if self.iteration % 200 == 0:
|
|
# img to cpu
|
|
t = 0
|
|
gt_local_frames_cpu = ((gt_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
|
masked_local_frames = ((masked_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
|
prop_local_frames_cpu = ((prop_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
|
pred_local_frames_cpu = ((pred_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
|
img_results = torch.cat([masked_local_frames[0][t], gt_local_frames_cpu[0][t],
|
|
prop_local_frames_cpu[0][t], pred_local_frames_cpu[0][t]], 1)
|
|
img_results = torchvision.utils.make_grid(img_results, nrow=1, normalize=True)
|
|
if self.gen_writer is not None:
|
|
self.gen_writer.add_image(f'img/img:inp-gt-res-{t}', img_results, self.iteration)
|
|
|
|
t = 5
|
|
if masked_local_frames.shape[1] > 5:
|
|
img_results = torch.cat([masked_local_frames[0][t], gt_local_frames_cpu[0][t],
|
|
prop_local_frames_cpu[0][t], pred_local_frames_cpu[0][t]], 1)
|
|
img_results = torchvision.utils.make_grid(img_results, nrow=1, normalize=True)
|
|
if self.gen_writer is not None:
|
|
self.gen_writer.add_image(f'img/img:inp-gt-res-{t}', img_results, self.iteration)
|
|
|
|
# flow to cpu
|
|
gt_flows_forward_cpu = flow_to_image(gt_flows_bi[0][0]).cpu()
|
|
masked_flows_forward_cpu = (gt_flows_forward_cpu[0] * (1-local_masks[0][0].cpu())).to(gt_flows_forward_cpu)
|
|
pred_flows_forward_cpu = flow_to_image(pred_flows_bi[0][0]).cpu()
|
|
|
|
flow_results = torch.cat([gt_flows_forward_cpu[0], masked_flows_forward_cpu, pred_flows_forward_cpu[0]], 1)
|
|
if self.gen_writer is not None:
|
|
self.gen_writer.add_image('img/flow:gt-pred', flow_results, self.iteration)
|
|
|
|
# console logs
|
|
if self.config['global_rank'] == 0:
|
|
pbar.update(1)
|
|
if not self.config['model']['no_dis']:
|
|
pbar.set_description((f"d: {dis_loss.item():.3f}; "
|
|
f"hole: {hole_loss.item():.3f}; "
|
|
f"valid: {valid_loss.item():.3f}"))
|
|
else:
|
|
pbar.set_description((f"hole: {hole_loss.item():.3f}; "
|
|
f"valid: {valid_loss.item():.3f}"))
|
|
|
|
if self.iteration % self.train_args['log_freq'] == 0:
|
|
if not self.config['model']['no_dis']:
|
|
logging.info(f"[Iter {self.iteration}] "
|
|
f"d: {dis_loss.item():.4f}; "
|
|
f"hole: {hole_loss.item():.4f}; "
|
|
f"valid: {valid_loss.item():.4f}")
|
|
else:
|
|
logging.info(f"[Iter {self.iteration}] "
|
|
f"hole: {hole_loss.item():.4f}; "
|
|
f"valid: {valid_loss.item():.4f}")
|
|
|
|
# saving models
|
|
if self.iteration % self.train_args['save_freq'] == 0:
|
|
self.save(int(self.iteration))
|
|
|
|
if self.iteration > self.train_args['iterations']:
|
|
break
|
|
|
|
train_data = self.prefetcher.next() |