GPUS = 4 phi = network_wrappers.FunctionFromVectorField( networks.tallUNet(unet=networks.UNet2ChunkyMiddle, dimension=3) ) psi = network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3)) pretrained_lowres_net = inverseConsistentNet.InverseConsistentNet( network_wrappers.DoubleNet(phi, psi), lambda x, y: torch.mean((x - y) ** 2), 100, ) network_wrappers.assignIdentityMap(pretrained_lowres_net, input_shape) network_wrappers.adjust_batch_size(pretrained_lowres_net, 12) trained_weights = torch.load( "results/dd_l400_continue_rescalegrad2/knee_aligner_resi_net1800" ) # trained_weights = torch.load("../results/dd_knee_l400_continue_smallbatch2/knee_aligner_resi_net9300") # rained_weights = torch.load("../results/double_deformable_knee3/knee_aligner_resi_net22200") pretrained_lowres_net.load_state_dict(trained_weights) hires_net = inverseConsistentNet.InverseConsistentNet( network_wrappers.DoubleNet( network_wrappers.DownsampleNet(pretrained_lowres_net.regis_net, dimension=3), network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3)),
torch.manual_seed(1) torch.cuda.manual_seed(1) np.random.seed(1) print("=" * 50) net = inverseConsistentNet.InverseConsistentNet( network_wrappers.DoubleNet( network_wrappers.RandomShift(0.25), network_wrappers.FunctionFromVectorField( networks.tallUNet2(dimension=2)), ), lambda x, y: torch.mean((x - y)**2), lmbda, ) input_shape = next(iter(d1))[0].size() network_wrappers.assignIdentityMap(net, input_shape) net.cuda() optimizer = torch.optim.Adam(net.parameters(), lr=0.001) net.train() xs = [] for _ in range(40): y = np.array(train.train2d(net, optimizer, d1, d2, epochs=50)) xs.append(y) x = np.concatenate(xs) plt.title("Loss curve for " + type(net.regis_net).__name__ + " lambda=" + str(lmbda)) plt.plot(x[:, :3]) plt.savefig(describe.run_dir + f"loss.png") plt.clf() plt.title("Log # pixels with negative Jacobian per epoch")
data_size=data_size, hollow=True, batch_size=batch_size) d1_t, d2_t = data.get_dataset_triangles("test", data_size=data_size, hollow=True, batch_size=batch_size) image_A, image_B = (x[0].cuda() for x in next(zip(d1, d2))) net = inverseConsistentNet.InverseConsistentNet( network_wrappers.FunctionFromMatrix(networks.ConvolutionalMatrixNet()), lambda x, y: torch.mean((x - y)**2), 100, ) network_wrappers.assignIdentityMap(net, image_A.shape) net.cuda() import train optim = torch.optim.Adam(net.parameters(), lr=0.00001) net.train().cuda() xs = [] for _ in range(240): y = np.array(train.train2d(net, optim, d1, d2, epochs=50)) xs.append(y) x = np.concatenate(xs) plt.title("Loss curve for " + type(net.regis_net).__name__) plt.plot(x[:, :3]) plt.savefig(describe.run_dir + f"loss.png")