class Epoch: def __init__(self, model, num_classes, losses, stage_name, device='cpu', verbose=True, EvaluatorIn=None): self.model = model self.losses = losses self.stage_name = stage_name self.verbose = verbose self.device = device self._to_device() if EvaluatorIn is None: self.evaluator = Evaluator(num_classes, device) else: self.evaluator = EvaluatorIn def _to_device(self): for loss in self.losses: loss.to(self.device) self.model.to(self.device) @classmethod def _format_logs(self, logs): str_logs = ['{} - {:.4}'.format(k, v) for k, v in logs.items()] s = ', '.join(str_logs) return s def batch_update(self, x, y): raise NotImplementedError def on_epoch_start(self): pass def run(self, dataloader): self.on_epoch_start() logs = {} loss_meter = AverageValueMeter() losses_meters = { loss.__name__: AverageValueMeter() for loss in self.losses } with tqdm(dataloader, desc=self.stage_name, file=sys.stdout, disable=not (self.verbose)) as iterator: for sample in iterator: x = sample['image'] y = sample['label'] if 'image_class' in sample: image_class = sample['image_class'] else: image_class = None x, y = x.to(self.device), y.to(self.device) loss, losses, y_pred = self.batch_update( x, y, image_class, sample) # update losses logs loss_value = loss.cpu().detach().numpy() loss_meter.add(loss_value) loss_logs = {'total_loss': loss_meter.mean} logs.update(loss_logs) for loss_fn in self.losses: loss_value = loss_fn(y_pred, y, x, image_class) if type(loss_value) == torch.Tensor: loss_value = loss_value.cpu().detach().numpy() losses_meters[loss_fn.__name__].add(loss_value) losses_logs = {k: v.mean for k, v in losses_meters.items()} logs.update(losses_logs) if self.verbose: s = self._format_logs(logs) iterator.set_postfix_str(s) logs['acc'] = self.evaluator.Pixel_Accuracy() logs['fscore'] = self.evaluator.f_score() logs['jaccard'] = self.evaluator.Jaccard() logs['miou'] = self.evaluator.Mean_Intersection_over_Union() logs['acc_class'] = self.evaluator.Pixel_Accuracy_Class() return logs
class Trainer(object): def __init__(self, args): self.args = args # Define Saver self.saver = Saver(args) self.saver.save_experiment_config() # Define Tensorboard Summary self.summary = TensorboardSummary(args.logdir) self.writer = self.summary.create_summary() # Define Dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} dltrain = DLDataset('trainval', "./data/pascal_voc_seg/tfrecord/") dlval = DLDataset('val', "./data/pascal_voc_seg/tfrecord/") # dltrain = DLDataset('trainval', "./data/pascal_voc_seg/VOCdevkit/VOC2012/") # dlval = DLDataset('val', "./data/pascal_voc_seg/VOCdevkit/VOC2012/") self.train_loader = DataLoader(dltrain, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) self.val_loader = DataLoader(dlval, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) # Define network model = Deeplab() train_params = [{ 'params': model.get_1x_lr_params(), 'lr': args.lr }, { 'params': model.get_10x_lr_params(), 'lr': args.lr * 10 }] # Define Optimizer optimizer = torch.optim.SGD(train_params, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) # Define Criterion # whether to use class balanced weights self.criterion = nn.CrossEntropyLoss(ignore_index=255).cuda() self.model, self.optimizer = model, optimizer # Define Evaluator self.evaluator = Evaluator(21) # Define lr scheduler self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer=optimizer) # Using cuda # if args.cuda: # self.model = torch.nn.DataParallel(self.model) self.model = self.model.cuda() # Resuming checkpoint self.best_pred = 0.0 if args.resume is not None: if not os.path.isfile(args.resume): raise RuntimeError("=> no checkpoint found at '{}'".format( args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] if args.cuda: self.model.module.load_state_dict(checkpoint['state_dict']) else: self.model.load_state_dict(checkpoint['state_dict']) if not args.ft: self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # Clear start epoch if fine-tuning if args.ft: args.start_epoch = 0 def training(self, epoch): train_loss = 0.0 self.model.train() tbar = tqdm(self.train_loader) num_img_tr = len(self.train_loader) for i, (image, target) in enumerate(tbar): if self.args.cuda: image, target = image.cuda(), target.cuda() self.optimizer.zero_grad() output = self.model(image) loss = self.criterion(output, target.long()) loss.backward() self.optimizer.step() train_loss += loss.item() tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1))) self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch) # Show 10 * 3 inference results each epoch # if i % (num_img_tr // 10) == 0: if i % 10 == 0: global_step = i + num_img_tr * epoch self.summary.visualize_image(self.writer, self.args.dataset, image, target, output, global_step) self.scheduler.step(train_loss) self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch) print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0])) print('Loss: %.3f' % train_loss) if self.args.no_val: # save checkpoint every epoch is_best = False self.saver.save_checkpoint( { 'epoch': epoch + 1, 'state_dict': self.model.module.state_dict(), 'optimizer': self.optimizer.state_dict(), 'best_pred': self.best_pred, }, is_best) def validation(self, epoch): self.model.eval() self.evaluator.reset() tbar = tqdm(self.val_loader, desc='\r') test_loss = 0.0 for i, (image, target) in enumerate(tbar): if self.args.cuda: image, target = image.cuda(), target.cuda() with torch.no_grad(): output = self.model(image) loss = self.criterion(output, target.long()) test_loss += loss.item() tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1))) pred = output.data.cpu().numpy() target = target.cpu().numpy() pred = np.argmax(pred, axis=1) # Add batch sample into evaluator self.evaluator.add_batch(target, pred) # Fast test during the training Acc = self.evaluator.Pixel_Accuracy() Acc_class = self.evaluator.Pixel_Accuracy_Class() mIoU = self.evaluator.Mean_Intersection_over_Union() FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union() self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch) self.writer.add_scalar('val/mIoU', mIoU, epoch) self.writer.add_scalar('val/Acc', Acc, epoch) self.writer.add_scalar('val/Acc_class', Acc_class, epoch) self.writer.add_scalar('val/fwIoU', FWIoU, epoch) print('Validation:') print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0])) print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format( Acc, Acc_class, mIoU, FWIoU)) print('Loss: %.3f' % test_loss) new_pred = mIoU if new_pred > self.best_pred: is_best = True self.best_pred = new_pred self.saver.save_checkpoint( { 'epoch': epoch + 1, 'state_dict': self.model.state_dict(), 'optimizer': self.optimizer.state_dict(), 'best_pred': self.best_pred, }, is_best)