validation function
(test_loader, model, epoch, save_path, writer)
| 185 | |
| 186 | |
| 187 | def val(test_loader, model, epoch, save_path, writer): |
| 188 | """ |
| 189 | validation function |
| 190 | """ |
| 191 | global best_mae, best_epoch |
| 192 | model.eval() |
| 193 | with torch.no_grad(): |
| 194 | mae_sum = 0 |
| 195 | # mae_sum_edge = 0 |
| 196 | for i in range(test_loader.size): |
| 197 | image, gt, name, img_for_post = test_loader.load_data() |
| 198 | gt = np.asarray(gt, np.float32) |
| 199 | gt /= (gt.max() + 1e-8) |
| 200 | image = image.cuda(device=device_ids[0]) |
| 201 | |
| 202 | result = model(image) |
| 203 | |
| 204 | res = F.upsample(result[4], size=gt.shape, mode='bilinear', align_corners=False) |
| 205 | res = res.sigmoid().data.cpu().numpy().squeeze() |
| 206 | res = (res - res.min()) / (res.max() - res.min() + 1e-8) |
| 207 | mae_sum += np.sum(np.abs(res - gt)) * 1.0 / (gt.shape[0] * gt.shape[1]) |
| 208 | |
| 209 | mae = mae_sum / test_loader.size |
| 210 | writer.add_scalar('MAE', torch.tensor(mae), global_step=epoch) |
| 211 | print('Epoch: {}, MAE: {}, bestMAE: {}, bestEpoch: {}.'.format(epoch, mae, best_mae, best_epoch)) |
| 212 | if epoch == 1: |
| 213 | best_mae = mae |
| 214 | best_epoch = 1 |
| 215 | else: |
| 216 | if mae < best_mae: |
| 217 | best_mae = mae |
| 218 | best_epoch = epoch |
| 219 | torch.save(model.state_dict(), save_path + 'Net_epoch_best.pth') |
| 220 | print('Save state_dict successfully! Best epoch:{}.'.format(epoch)) |
| 221 | logging.info( |
| 222 | '[Val Info]:Epoch:{} MAE:{} bestEpoch:{} bestMAE:{}'.format(epoch, mae, best_epoch, best_mae)) |
| 223 | |
| 224 | |
| 225 | if __name__ == '__main__': |