def __init__(self, config, model, train_loader, val_loader=None):
super(Trainer, self).__init__(config, model, train_loader, val_loader)
if self.device != torch.device('cpu'):
self.train_loader = DataPrefetcher(train_loader, device=self.device)
if self.val:
self.val_loader = DataPrefetcher(val_loader, device=self.device)
def __init__(self,
model,
loss,
resume,
config,
train_loader,
val_loader=None,
train_logger=None,
prefetch=True):
""" Trainer 类
__init__:
1、TRANSORMS FOR VISUALIZATION
2、预读取
_train_epoch：
"""
super(Trainer, self).__init__(model, loss, resume, config,
train_loader, val_loader, train_logger)
self.wrt_mode, self.wrt_step = 'train_', 0
self.log_step = config['trainer'].get(
'log_per_iter', int(np.sqrt(self.train_loader.batch_size)))
if config['trainer']['log_per_iter']:
self.log_step = int(
self.log_step / self.train_loader.batch_size) + 1
self.num_classes = self.train_loader.dataset.num_classes
# TRANSORMS FOR VISUALIZATION
self.restore_transform = transforms.Compose([
local_transforms.DeNormalize(self.train_loader.MEAN,
self.train_loader.STD),
transforms.ToPILImage()
])
self.viz_transform = transforms.Compose(
[transforms.Resize((400, 400)),
transforms.ToTensor()])
# 预读取
if self.device == torch.device('cpu'):
prefetch = False
if prefetch:
self.train_loader = DataPrefetcher(train_loader,
device=self.device)
self.val_loader = DataPrefetcher(val_loader, device=self.device)
torch.backends.cudnn.benchmark = True