def check_split_and_input_ch(self, split, input_ch):
modal = split.split("_")[-1]
if modal == "rgbd":
assert input_ch == 4
print ("4ch is Depth channel")
elif modal == "d":
assert input_ch == 1
if modal == "rgb" and input_ch == 4:
print (emphasize_str("4ch is R channel"))
if torch.cuda.is_available():
model_g1.cuda()
model_g2.cuda()
model_f1.cuda()
weight = weight.cuda()
criterion = CrossEntropyLoss2d(weight)
configure(args.tflog_dir, flush_secs=5)
model_g1.train()
model_g2.train()
model_f1.train()
if args.fix_bn:
print(emphasize_str("BN layers are NOT trained!"))
fix_batchnorm_when_training(model_g1)
fix_batchnorm_when_training(model_g2)
fix_batchnorm_when_training(model_f1)
# check_training(model)
for epoch in range(start_epoch, args.epochs):
epoch_loss = 0
for ind, (images, labels) in tqdm(enumerate(train_loader)):
imgs = Variable(images)
lbls = Variable(labels)
if torch.cuda.is_available():
imgs, lbls = imgs.cuda(), lbls.cuda()
