def __init__(self, in_channel, n_classes, start_channel, is_train=True, imgshape=(160, 192, 144), range_flow=0.4,
model_lvl2=None):
super(Miccai2020_LDR_laplacian_unit_disp_add_lvl3, self).__init__()
self.in_channel = in_channel
self.n_classes = n_classes
self.start_channel = start_channel
self.range_flow = range_flow
self.is_train = is_train
self.imgshape = imgshape
self.model_lvl2 = model_lvl2
self.grid_1 = generate_grid_unit(self.imgshape)
self.grid_1 = torch.from_numpy(np.reshape(self.grid_1, (1,) + self.grid_1.shape)).cuda().float()
self.transform = SpatialTransform_unit().cuda()
bias_opt = False
self.input_encoder_lvl1 = self.input_feature_extract(self.in_channel+3, self.start_channel * 4, bias=bias_opt)
self.down_conv = nn.Conv3d(self.start_channel * 4, self.start_channel * 4, 3, stride=2, padding=1, bias=bias_opt)
self.resblock_group_lvl1 = self.resblock_seq(self.start_channel * 4, bias_opt=bias_opt)
self.up_tri = torch.nn.Upsample(scale_factor=2, mode="trilinear")
self.up = nn.ConvTranspose3d(self.start_channel * 4, self.start_channel * 4, 2, stride=2,
padding=0, output_padding=0, bias=bias_opt)
self.output_lvl1 = self.outputs(self.start_channel * 8, self.n_classes, kernel_size=3, stride=1, padding=1, bias=False)
def __init__(self, in_channel, n_classes, start_channel, is_train=True, imgshape=(160, 192, 144), range_flow=0.4):
super(Miccai2020_LDR_laplacian_unit_add_lvl1, self).__init__()
self.in_channel = in_channel
self.n_classes = n_classes
self.start_channel = start_channel
self.range_flow = range_flow
self.is_train = is_train
self.imgshape = imgshape
self.grid_1 = generate_grid_unit(self.imgshape)
self.grid_1 = torch.from_numpy(np.reshape(self.grid_1, (1,) + self.grid_1.shape)).cuda().float()
self.diff_transform = DiffeomorphicTransform_unit(time_step=7).cuda()
self.transform = SpatialTransform_unit().cuda()
bias_opt = False
self.input_encoder_lvl1 = self.input_feature_extract(self.in_channel, self.start_channel * 4, bias=bias_opt)
self.down_conv = nn.Conv3d(self.start_channel * 4, self.start_channel * 4, 3, stride=2, padding=1, bias=bias_opt)
self.resblock_group_lvl1 = self.resblock_seq(self.start_channel * 4, bias_opt=bias_opt)
self.up = nn.ConvTranspose3d(self.start_channel * 4, self.start_channel * 4, 2, stride=2,
padding=0, output_padding=0, bias=bias_opt)
self.down_avg = nn.AvgPool3d(kernel_size=3, stride=2, padding=1, count_include_pad=False)
self.output_lvl1 = self.outputs(self.start_channel * 8, self.n_classes, kernel_size=3, stride=1, padding=1, bias=False)
def test():
imgshape_4 = (160 / 4, 192 / 4, 144 / 4)
imgshape_2 = (160 / 2, 192 / 2, 144 / 2)
model_lvl1 = Miccai2020_LDR_laplacian_unit_disp_add_lvl1(2, 3, start_channel, is_train=True, imgshape=imgshape_4,
range_flow=range_flow).cuda()
model_lvl2 = Miccai2020_LDR_laplacian_unit_disp_add_lvl2(2, 3, start_channel, is_train=True, imgshape=imgshape_2,
range_flow=range_flow, model_lvl1=model_lvl1).cuda()
model = Miccai2020_LDR_laplacian_unit_disp_add_lvl3(2, 3, start_channel, is_train=False, imgshape=imgshape,
range_flow=range_flow, model_lvl2=model_lvl2).cuda()
transform = SpatialTransform_unit().cuda()
model.load_state_dict(torch.load(opt.modelpath))
model.eval()
transform.eval()
grid = generate_grid_unit(imgshape)
grid = torch.from_numpy(np.reshape(grid, (1,) + grid.shape)).cuda().float()
use_cuda = True
device = torch.device("cuda" if use_cuda else "cpu")
fixed_img = load_4D(fixed_path)
moving_img = load_4D(moving_path)
fixed_img = torch.from_numpy(fixed_img).float().to(device).unsqueeze(dim=0)
moving_img = torch.from_numpy(moving_img).float().to(device).unsqueeze(dim=0)
with torch.no_grad():
F_X_Y = model(moving_img, fixed_img)
X_Y = transform(moving_img, F_X_Y.permute(0, 2, 3, 4, 1), grid).data.cpu().numpy()[0, 0, :, :, :]
F_X_Y_cpu = F_X_Y.data.cpu().numpy()[0, :, :, :, :].transpose(1, 2, 3, 0)
F_X_Y_cpu = transform_unit_flow_to_flow(F_X_Y_cpu)
save_flow(F_X_Y_cpu, savepath+'/warpped_flow.nii.gz')
save_img(X_Y, savepath+'/warpped_moving.nii.gz')
print("Finished")
def train_lvl3():
print("Training lvl3...")
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model_lvl1 = Miccai2020_LDR_laplacian_unit_disp_add_lvl1(
2,
3,
start_channel,
is_train=True,
imgshape=imgshape_4,
range_flow=range_flow).to(device)
model_lvl2 = Miccai2020_LDR_laplacian_unit_disp_add_lvl2(
2,
3,
start_channel,
is_train=True,
imgshape=imgshape_2,
range_flow=range_flow,
model_lvl1=model_lvl1).to(device)
model_path = sorted(
glob.glob("../Model/Stage/" + model_name + "stagelvl2_?????.pth"))[-1]
model_lvl2.load_state_dict(torch.load(model_path))
print("Loading weight for model_lvl2...", model_path)
# Freeze model_lvl1 weight
for param in model_lvl2.parameters():
param.requires_grad = False
model = Miccai2020_LDR_laplacian_unit_disp_add_lvl3(
2,
3,
start_channel,
is_train=True,
imgshape=imgshape,
range_flow=range_flow,
model_lvl2=model_lvl2).to(device)
loss_similarity = multi_resolution_NCC(win=7, scale=3)
loss_smooth = smoothloss
loss_Jdet = neg_Jdet_loss
transform = SpatialTransform_unit().to(device)
transform_nearest = SpatialTransformNearest_unit().to(device)
for param in transform.parameters():
param.requires_grad = False
param.volatile = True
# OASIS
names = sorted(glob.glob(datapath + '/*.nii'))
grid = generate_grid(imgshape)
grid = torch.from_numpy(np.reshape(grid,
(1, ) + grid.shape)).to(device).float()
grid_unit = generate_grid_unit(imgshape)
grid_unit = torch.from_numpy(np.reshape(
grid_unit, (1, ) + grid_unit.shape)).to(device).float()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
model_dir = '../Model'
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
lossall = np.zeros((4, iteration_lvl3 + 1))
training_generator = Data.DataLoader(Dataset_epoch(names, norm=False),
batch_size=1,
shuffle=True,
num_workers=2)
step = 0
load_model = False
if load_model is True:
model_path = "../Model/LDR_LPBA_NCC_lap_share_preact_1_05_3000.pth"
print("Loading weight: ", model_path)
step = 3000
model.load_state_dict(torch.load(model_path))
temp_lossall = np.load(
"../Model/loss_LDR_LPBA_NCC_lap_share_preact_1_05_3000.npy")
lossall[:, 0:3000] = temp_lossall[:, 0:3000]
while step <= iteration_lvl3:
for X, Y in training_generator:
X = X.to(device).float()
Y = Y.to(device).float()
# compose_field_e0_lvl1, warpped_inputx_lvl1_out, y, output_disp_e0_v, lvl1_v, lvl2_v, e0
F_X_Y, X_Y, Y_4x, F_xy, F_xy_lvl1, F_xy_lvl2, _ = model(X, Y)
# 3 level deep supervision NCC
loss_multiNCC = loss_similarity(X_Y, Y_4x)
F_X_Y_norm = transform_unit_flow_to_flow_cuda(
F_X_Y.permute(0, 2, 3, 4, 1).clone())
loss_Jacobian = loss_Jdet(F_X_Y_norm, grid)
# reg2 - use velocity
_, _, x, y, z = F_X_Y.shape
F_X_Y[:, 0, :, :, :] = F_X_Y[:, 0, :, :, :] * z
F_X_Y[:, 1, :, :, :] = F_X_Y[:, 1, :, :, :] * y
F_X_Y[:, 2, :, :, :] = F_X_Y[:, 2, :, :, :] * x
loss_regulation = loss_smooth(F_X_Y)
loss = loss_multiNCC + antifold * loss_Jacobian + smooth * loss_regulation
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # backpropagation, compute gradients
optimizer.step() # apply gradients
lossall[:, step] = np.array([
loss.item(),
loss_multiNCC.item(),
loss_Jacobian.item(),
loss_regulation.item()
])
sys.stdout.write(
"\r" +
'step "{0}" -> training loss "{1:.4f}" - sim_NCC "{2:4f}" - Jdet "{3:.10f}" -smo "{4:.4f}"'
.format(step, loss.item(), loss_multiNCC.item(),
loss_Jacobian.item(), loss_regulation.item()))
sys.stdout.flush()
# with lr 1e-3 + with bias
if (step % n_checkpoint == 0):
modelname = model_dir + '/' + model_name + "stagelvl3_" + str(
step) + '.pth'
torch.save(model.state_dict(), modelname)
np.save(
model_dir + '/loss' + model_name + "stagelvl3_" +
str(step) + '.npy', lossall)
# Validation
if step == freeze_step:
model.unfreeze_modellvl2()
step += 1
if step > iteration_lvl3:
break
print("one epoch pass")
np.save(model_dir + '/loss' + model_name + 'stagelvl3.npy', lossall)