labeller.cuda()
optimiser = optim.Adam(ds_unet.parameters(), lr=learning_rate)
print('Project name ', project_name)
train_dices = []
train_losses = []
val_a_dices = []
val_a_losses = []
val_b_dices = []
val_b_losses = []
for i in range(epochs):
ds_unet.set_domain(DOMAIN_A)
train_dice, train_loss = train_segmentation_net(ds_unet,
train_a_loader,
'A',
0,
1,
criterion_seg,
opt=optimiser,
is_train=True,
name='dsbn a train',
epoch=i)
val_dice, val_loss = train_segmentation_net(ds_unet,
val_a_loader,
'A',
0,
loss_seg_a.backward()
optimiser_ps.step()
else:
pseudo.eval()
image_a = data['A'][0].cuda()
target_a = data['A'][1].cuda()
image_b = data['B'][0].cuda()
edges_b = data['B'][2].cuda()
pseudo_b = pseudo.downsample(edges_b)
pred_pseudo_b = pseudo.upsample(*pseudo_b)
# pred_pseudo_b = pseudo(edges_b)
target_b = torch.round(pred_pseudo_b).detach().cuda()
net.set_domain(DOMAIN_A)
res_a = net.downsample(image_a)
pred_seg_a = net.upsample(*res_a)
net.set_domain(DOMAIN_B)
res_b = net.downsample(image_b)
pred_seg_b = net.upsample(*res_b)
loss_seg_a = criterion(pred_seg_a, target_a)
loss_seg_b = criterion(pred_seg_b, target_b)
loss = loss_seg_a + loss_seg_b
loss.backward()
optimiser_net.step()
