()
| 225 | |
| 226 | |
| 227 | def visualize_stn(): |
| 228 | with torch.no_grad(): |
| 229 | # Get a batch of training data |
| 230 | data = next(iter(test_loader))[0].to(device) |
| 231 | |
| 232 | input_tensor = data.cpu() |
| 233 | transformed_input_tensor = model.stn(data).cpu() |
| 234 | |
| 235 | in_grid = convert_image_np( |
| 236 | torchvision.utils.make_grid(input_tensor)) |
| 237 | |
| 238 | out_grid = convert_image_np( |
| 239 | torchvision.utils.make_grid(transformed_input_tensor)) |
| 240 | |
| 241 | # Plot the results side-by-side |
| 242 | f, axarr = plt.subplots(1, 2) |
| 243 | axarr[0].imshow(in_grid) |
| 244 | axarr[0].set_title('Dataset Images') |
| 245 | |
| 246 | axarr[1].imshow(out_grid) |
| 247 | axarr[1].set_title('Transformed Images') |
| 248 | |
| 249 | for epoch in range(1, 20 + 1): |
| 250 | train(epoch) |
no test coverage detected