| 9 | import numpy as np |
| 10 | |
| 11 | def plot_correlation(forces_gt, forces_pred, log_path, epoch): |
| 12 | colors = ["#7998e8", "#52a375", "#803b6b"] |
| 13 | |
| 14 | correlation_fig = plt.figure(figsize=(20, 5)) |
| 15 | axs: np.ndarray = correlation_fig.subplots(1, 3) |
| 16 | for i, (force_gt, force_pred) in enumerate(zip(forces_gt.T, forces_pred.T)): |
| 17 | axs[i].scatter( |
| 18 | force_gt, |
| 19 | force_pred, |
| 20 | s=2, |
| 21 | color=colors[i] |
| 22 | ) |
| 23 | axs[i].set_xlabel("Ground Truth (N)") |
| 24 | axs[i].set_ylabel("Prediction (N)") |
| 25 | axs[i].set_title(f"Force {['X', 'Y', 'Z'][i]}") |
| 26 | axs[i].grid(True) |
| 27 | # plot 1:1 line |
| 28 | axs[i].plot( |
| 29 | [force_gt.min(), force_gt.max()], |
| 30 | [force_gt.min(), force_gt.max()], |
| 31 | "--", |
| 32 | color="gray", |
| 33 | ) |
| 34 | axs[i].legend() |
| 35 | # return correlation_fig, axs |
| 36 | |
| 37 | plt.savefig(log_path+'/correlation_epoch'+str(epoch)+'.png') |
| 38 | plt.close("all") |
| 39 | |
| 40 | def train_one_epoch(model: torch.nn.Module, |
| 41 | data_loader, optimizer: torch.optim.Optimizer, |