def test_PEM(data_loader, model, epoch, writer, opt): model.eval() epoch_iou_loss = 0 losses = AverageMeter() for n_iter, (input_data, label_iou) in enumerate(data_loader): PEM_output = model(input_data) iou_loss = PEM_loss_function(PEM_output, label_iou, model, opt) epoch_iou_loss += iou_loss.cpu().detach().numpy() losses.update(iou_loss.item()) if (n_iter + 1) % opt['print_freq'] == 0: print('[TEST] Epoch {}, iter {} / {}, loss: {}'.format( epoch, n_iter + 1, len(data_loader), losses.avg)) writer.add_scalars('data/iou_loss', {'validation': epoch_iou_loss / (n_iter + 1)}, epoch) print("PEM testing loss(epoch %d): iou - %.04f" % (epoch, epoch_iou_loss / (n_iter + 1))) state = {'epoch': epoch + 1, 'state_dict': model.state_dict()} torch.save(state, opt["checkpoint_path"] + "/pem_checkpoint.pth.tar") if epoch_iou_loss < model.module.pem_best_loss: model.module.pem_best_loss = np.mean(epoch_iou_loss) torch.save(state, opt["checkpoint_path"] + "/pem_best.pth.tar")
def train_PEM(data_loader, model, optimizer, epoch, writer, opt): model.train() epoch_iou_loss = 0 losses = AverageMeter() for n_iter, (input_data, label_iou) in enumerate(data_loader): PEM_output = model(input_data) iou_loss = PEM_loss_function(PEM_output, label_iou, model, opt) optimizer.zero_grad() iou_loss.backward() optimizer.step() epoch_iou_loss += iou_loss.cpu().detach().numpy() losses.update(iou_loss.item()) if (n_iter + 1) % opt['print_freq'] == 0: print('[TRAIN] Epoch {}, iter {} / {}, loss: {}'.format( epoch, n_iter + 1, len(data_loader), losses.avg)) writer.add_scalars('data/iou_loss', {'train': epoch_iou_loss / (n_iter + 1)}, epoch) print("PEM training loss(epoch %d): iou - %.04f" % (epoch, epoch_iou_loss / (n_iter + 1)))