def validation(args, val_loader, network, epoch, loss, optimizer, writer, saver, summary, best_pred, classnames): # Define Evaluator evaluator = Evaluator(args.num_classes) test_loss = 0.0 evaluator.reset() #Put the model in eval() mode network.eval() network.cuda() num_img_tr = len(val_loader) for i, sample in enumerate(val_loader): # Get an output from the network img, target = sample with torch.no_grad(): img = img.to(device) target = target.to(device) output = network(img) test_loss += loss(output, target) pred = output.data.cpu().numpy() target = target.cpu().numpy() pred = np.argmax(pred, axis=1) # Add batch sample into evaluator evaluator.add_batch(target, pred) if i % (num_img_tr // 1) == 0: global_step = i + num_img_tr * epoch summary.visualize_image(writer, img, target, pred, global_step) # sample, img, target, output, prediction = None, None, None, None, None val_loader = None #val_loader = None # Fast test during the training Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() F1score = evaluator.F1_Score() Recall = evaluator.Recall() Precision = evaluator.Precision() IoU = evaluator.IoU() evaluator.plot_confusion_matrix(evaluator.confusion_matrix, classes=classnames, normalize=True) # Send the confusion matrix to TensorBoard confusion_matrix = plt.imread('/home/student/aeroscan/tileimp/test.png')[:,:,:3] confusion_matrix = np.transpose(confusion_matrix, (2,0,1)) confusion_matrix = torch.from_numpy(confusion_matrix) writer.add_image('confusion matrix', confusion_matrix, global_step) writer.add_scalar('val/total_loss_epoch', test_loss, epoch) writer.add_scalar('val/mIoU', mIoU, epoch) writer.add_scalar('val/Acc', Acc, epoch) writer.add_scalar('val/Acc_class', Acc_class, epoch) writer.add_scalar('val/fwIoU', FWIoU, epoch) writer.add_scalar("val/F1score", F1score, epoch) writer.add_scalar("val/Recall", Recall, epoch) writer.add_scalar("val/Precision", Precision, epoch) writer.add_scalar("val/IoU", Precision, epoch) print('Validation:') print('[Epoch:', epoch) print(""" Acc:{} Acc_class:{} mIoU:{} fwIoU:{} Precision:{} F1score:{} Recall:{} IoU:{}""" .format(Acc,Acc_class,mIoU,FWIoU,Precision,F1score,Recall,IoU)) print('Loss: %.3f' % test_loss) new_pred = Precision if new_pred > best_pred: print("New best pred.txt.") is_best = True best_pred = new_pred saver.save_checkpoint({ 'epoch': epoch + 1, 'state_dict': network.state_dict(), 'optimizer': optimizer.state_dict(), 'best_pred': best_pred, }, is_best) return best_pred, saver
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(self.saver.experiment_dir) self.writer = self.summary.create_summary() # Define Dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args, **kwargs) # Define network model = DFPENet(num_classes=self.nclass, backbone=args.backbone, output_stride=args.out_stride, sync_bn=args.sync_bn, freeze_bn=args.freeze_bn) 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 if args.use_balanced_weights: classes_weights_path = os.path.join(Path.db_root_dir(args.dataset), args.dataset+'_classes_weights.npy') if os.path.isfile(classes_weights_path): weight = np.load(classes_weights_path) else: weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass) weight = torch.from_numpy(weight.astype(np.float32)) else: weight = None self.criterion = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type) self.model, self.optimizer = model, optimizer # Define Evaluator self.evaluator = Evaluator(self.nclass) # Define lr scheduler self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.train_loader)) # Using cuda if args.cuda: self.model = self.model.cuda() self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids) patch_replication_callback(self.model) # 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.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, sample in enumerate(tbar): image, target = sample['image'], sample['label'] if self.args.cuda: image, target = image.cuda(), target.cuda() self.scheduler(self.optimizer, i, epoch, self.best_pred) self.optimizer.zero_grad() output = self.model(image) loss = self.criterion(output, target) 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 // 1) == 0: global_step = i + num_img_tr * epoch self.summary.visualize_image(self.writer, self.args.dataset, image, target, output, global_step) 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.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, sample in enumerate(tbar): image, target = sample['image'], sample['label'] if self.args.cuda: image, target = image.cuda(), target.cuda() with torch.no_grad(): output = self.model(image) loss = self.criterion(output, target) 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) if epoch==200: self.evaluator.plot_confusion_matrix(epoch) # Fast test during the training Rec = self.evaluator.Pixel_Accuracy_ALLClass() Pre = self.evaluator.Pixel_Precision_ALLClass() F1 = self.evaluator.F1_ALLClass() F1_mean = self.evaluator.F1_MEANClass() IoU = self.evaluator.Class_Intersection_over_Union() 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/Rec[1]', Rec[1], epoch) self.writer.add_scalar('val/Pre[1]', Pre[1], epoch) self.writer.add_scalar('val/F1[0]', F1[0], epoch) self.writer.add_scalar('val/F1[1]', F1[1], epoch) self.writer.add_scalar('val/F1_mean', F1_mean, epoch) self.writer.add_scalar('val/IoU[0]', IoU[0], epoch) self.writer.add_scalar('val/IoU[1]', IoU[1], 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("F1[0]:{}, F1[1]:{}, F1_mean: {}".format(F1[0], F1[1], F1_mean, )) print("IoU[0]:{}, IoU[1]:{}, mIoU: {}".format(IoU[0], IoU[1], mIoU)) print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU)) print("Rec[1]:{}, Pre[1]:{}".format(Rec[1], Pre[1])) print('Loss: %.3f' % test_loss) filename = "./rec.txt" with open(filename,'a', encoding='utf-8') as f: f.writelines(str(Rec[1])+'\n') filename1 = "./pre.txt" with open(filename1,'a', encoding='utf-8') as f1: f1.writelines(str(Pre[1])+'\n') filename2 = "./miou.txt" with open(filename2,'a', encoding='utf-8') as f2: f2.writelines(str(IoU[1])+'\n') new_pred = IoU[1] 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)