else:
# If in Supervision Tuning mode.
trainer.set_sup_trainable(True)
trainer.set_gen_trainable(False)
labels_1 = labels_1.to(dtype=torch.long)
labels_1[labels_1 > 0] = 1
labels_1 = Variable(labels_1.cuda(), requires_grad=False)
labels_2 = labels_2.to(dtype=torch.long)
labels_2[labels_2 > 0] = 1
labels_2 = Variable(labels_2.cuda(), requires_grad=False)
trainer.sup_update(images_1, images_2, labels_1, labels_2, index_1, index_2, use_1, use_2, config)
if (ep + 1) % config['snapshot_save_iter'] == 0:
trainer.save(checkpoint_directory, (ep + 1))
for i in range(config['n_datasets']):
print(' Testing ' + dataset_letters[i] + '...')
jacc_list = list()
for it, data in enumerate(test_loader_list[i]):
images = data[0]
labels = data[1]
use = data[2]
dis_loss+=trainer.dis_update(images_1, images_2, index_1, index_2, config)
gen_loss+=trainer.gen_update(images_1, images_2, index_1, index_2, config)
else:
# If in Supervision Tuning mode.
trainer.set_sup_trainable(True)
trainer.set_gen_trainable(False)
labels_1 = labels_1.to(dtype=torch.long)
labels_1 = Variable(labels_1.cuda(), requires_grad=False)
labels_2 = labels_2.to(dtype=torch.long)
labels_2 = Variable(labels_2.cuda(), requires_grad=False)
if (ep+1)<=10:
temp_loss=trainer.sup_update(images_1, images_2, labels_1, labels_2, index_1, index_2, use_1, use_2,ep, config)
seg_loss+=temp_loss[0]
seg_gen_loss+=temp_loss[1]
else:
temp_loss=trainer.sup_update(images_1, images_2, labels_1, labels_2, index_1, index_2, use_1, use_2,ep, config)
seg_loss+=temp_loss[0]
seg_gen_loss+=temp_loss[1]
dis2_loss+=trainer.dis2_update(images_1,images_2,index_1, index_2, use_1, use_2, config)
gen_loss=gen_loss/(it+1)
seg_loss=seg_loss/(it+1)
seg_gen_loss=seg_gen_loss/(it+1)
dis_loss=dis_loss/(it+1)
dis2_loss=dis2_loss/(it+1)
writer.add_scalar('train_seg/seg_loss', seg_loss, ep+1)
writer.add_scalar('train_seg2/seg_gen_loss', seg_gen_loss, ep+1)
