(dataset, model, model_input, gt, model_output, writer, total_steps, prefix='train')
| 25 | |
| 26 | |
| 27 | def write_simple_1D_function_summary(dataset, model, model_input, gt, model_output, writer, total_steps, prefix='train'): |
| 28 | |
| 29 | model_input, gt = dataset[0] |
| 30 | model_input = {k: v.unsqueeze(0) for k, v in model_input.items()} |
| 31 | model_input = training.dict2cuda(model_input) |
| 32 | model_output = model(model_input) |
| 33 | |
| 34 | pred_func = to_numpy(model_output['model_out']['output'].squeeze()) # B, Samples, DimOut |
| 35 | coords = to_numpy(gt['coords'].squeeze()) # B, Samples, DimIn |
| 36 | |
| 37 | val_coords = coords |
| 38 | val_pred_func = pred_func |
| 39 | val_gt_func = to_numpy(gt['func'].squeeze()) # B, Samples, DimOut |
| 40 | |
| 41 | idx = model_input['idx'].cpu().long().detach().numpy().squeeze() |
| 42 | train_coords = coords[idx] |
| 43 | train_pred_func = pred_func[idx] |
| 44 | |
| 45 | fig = plt.figure() |
| 46 | plt.subplot(211) |
| 47 | plt.plot(val_coords, val_gt_func, label='GT', linewidth=2) |
| 48 | plt.plot(val_coords, val_pred_func, label='Val') |
| 49 | plt.plot(train_coords, train_pred_func, '.', label='Train', markersize=8) |
| 50 | |
| 51 | plt.ylim([-1, 1]) |
| 52 | plt.legend() |
| 53 | |
| 54 | plt.subplot(212) |
| 55 | oversample = 32 |
| 56 | |
| 57 | freqs = np.fft.fftshift(np.fft.fftfreq(oversample * len(val_coords), 1/len(val_coords))) |
| 58 | val_gt_func = np.pad(val_gt_func, ((0, (oversample-1)*len(val_gt_func)))) |
| 59 | val_pred_func = np.pad(val_pred_func, ((0, (oversample-1)*len(val_pred_func)))) |
| 60 | |
| 61 | plt.plot(freqs, np.fft.fftshift(np.abs(np.fft.fft(val_gt_func))), label='GT', linewidth=2) |
| 62 | plt.plot(freqs, np.fft.fftshift(np.abs(np.fft.fft(val_pred_func))), label='Val') |
| 63 | plt.xlim(-32, 32) |
| 64 | |
| 65 | plt.tight_layout() |
| 66 | writer.add_figure(prefix + '/gt_vs_pred', fig, global_step=total_steps) |
| 67 | |
| 68 | |
| 69 | def write_multiscale_image_summary(image_resolution, train_dataset, model, model_input, gt, |
nothing calls this directly
no test coverage detected