def test_2d_registration_train(self): import icon_registration.data as data import icon_registration.networks as networks import icon_registration.network_wrappers as network_wrappers import icon_registration.train as train import icon_registration.inverseConsistentNet as inverseConsistentNet import numpy as np import torch import random import os random.seed(1) torch.manual_seed(1) torch.cuda.manual_seed(1) np.random.seed(1) batch_size = 128 d1, d2 = data.get_dataset_triangles("train", data_size=50, hollow=False, batch_size=batch_size) d1_t, d2_t = data.get_dataset_triangles("test", data_size=50, hollow=False, batch_size=batch_size) lmbda = 2048 print("=" * 50) net = inverseConsistentNet.InverseConsistentNet( network_wrappers.FunctionFromVectorField( networks.tallUNet2(dimension=2)), # Our image similarity metric. The last channel of x and y is whether the value is interpolated or extrapolated, # which is used by some metrics but not this one lambda x, y: torch.mean((x[:, :1] - y[:, :1])**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() y = np.array(train.train2d(net, optimizer, d1, d2, epochs=50)) # Test that image similarity is good enough self.assertLess(np.mean(y[-5:, 1]), 0.1) # Test that folds are rare enough self.assertLess(np.mean(np.exp(y[-5:, 3] - 0.1)), 2) print(y)
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_net25500" ) # 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.DownsampleNet(pretrained_lowres_net.regis_net, dimension=3), lambda x, y: torch.mean((x - y) ** 2), 800,
image_A, image_B = (x[0].cuda() for x in next(zip(d1, d2))) net = inverseConsistentNet.InverseConsistentNet( network_wrappers.DoubleNet( network_wrappers.FunctionFromVectorField( networks.tallUNet2(dimension=2)), network_wrappers.FunctionFromVectorField( networks.tallUNet2(dimension=2)), ), lambda x, y: torch.mean((x - y)**2), 700, ) input_shape = next(iter(d1))[0].size() network_wrappers.assignIdentityMap(net, input_shape) net.cuda() import icon_registration.train as train optim = torch.optim.Adam(net.parameters(), lr=0.0001) 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(footsteps.output_dir + f"loss.png")
data_size = 50 d1, d2 = data.get_dataset_triangles( "train", 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 icon_registration.train as 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])
psi = network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3)) pretrained_lowres_net = network_wrappers.DoubleNet(phi, psi) hires_net = inverseConsistentNet.InverseConsistentNet( network_wrappers.DoubleNet( network_wrappers.DownsampleNet(pretrained_lowres_net, dimension=3), network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3)), ), inverseConsistentNet.ssd_only_interpolated, 3600, ) BATCH_SIZE = 4 SCALE = 2 # 1 IS QUARTER RES, 2 IS HALF RES, 4 IS FULL RES input_shape = [BATCH_SIZE, 1, 40 * SCALE, 96 * SCALE, 96 * SCALE] network_wrappers.assignIdentityMap(hires_net, input_shape) trained_weights = torch.load("results/hires_smart_6/knee_aligner_resi_net74700") hires_net.load_state_dict(trained_weights) fourth_net = inverseConsistentNet.InverseConsistentNet( network_wrappers.DoubleNet( hires_net.regis_net, network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3)), ), inverseConsistentNet.ssd_only_interpolated, 3600, ) for p in fourth_net.regis_net.netPhi.parameters(): p.requires_grad = False