def train_base(generator, optimG, trainloader, valoader, args):
best_miou = -1
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
for batch_id, (img, mask, _) in enumerate(trainloader):
if args.nogpu:
img, mask = Variable(img), Variable(mask)
else:
img, mask = Variable(img.cuda()), Variable(mask.cuda())
itr = len(trainloader) * (epoch - 1) + batch_id
cprob = generator(img)
cprob = nn.LogSoftmax()(cprob)
Lseg = nn.NLLLoss2d()(cprob, mask)
poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.zero_grad()
Lseg.backward()
optimG.step()
print("[{}][{}]Loss: {:0.4f}".format(epoch, itr, Lseg.data[0]))
best_miou = snapshot(generator, valoader, epoch, best_miou,
args.snapshot_dir, args.prefix)
def train_base(generator, optimG, trainloader, valoader, f_path):
best_miou_unlabel = -1
best_miou = -1
for epoch in range(start_epoch, max_epoch + 1):
generator.train()
for batch_id, (img, mask, _, _) in enumerate(trainloader):
if not use_cuda:
img, mask = img, mask
else:
img, mask = img.cuda(), mask.cuda()
# print(mask.max())###测试mask的最大值是多少
itr = len(trainloader) * (epoch - 1) + batch_id
cprob = generator(img)
cprob = nn.LogSoftmax()(cprob)
# print(cprob.shape)
# print(mask.shape)
Lseg = nn.NLLLoss()(cprob, mask)
optimG = poly_lr_scheduler(optimG, g_lr, itr)
optimG.zero_grad()
Lseg.backward()
optimG.step()
print("[{}][{}]Loss: {:0.4f}".format(epoch, itr, Lseg.item()))
best_miou_unlabel, best_miou = snapshot(
generator, optimG, trainloader, valoader, epoch, best_miou_unlabel,
best_miou, snapshot_dir, prefix + '_' + str(epoch), f_path)
def train_semi(generator,discriminator,optimG,optimD,trainloader_l,trainloader_u,valoader,args):
best_miou = -1
best_eiou = -1
for epoch in range(args.start_epoch,args.max_epoch+1):
generator.train()
trainloader_l_iter = iter(trainloader_l)
trainloader_u_iter = iter(trainloader_u)
print("Epoch: {}".format(epoch))
batch_id = 0
while(True):
batch_id += 1
itr = (len(trainloader_u) + len(trainloader_l))*(epoch-1) + batch_id
LGsemi_d = 0
LGsemi_c = 0
if epoch > args.wait_semi:
if random.random() <0.5:
loader_l = trainloader_l_iter
loader_u = trainloader_u_iter
else:
loader_l = trainloader_u_iter
loader_u = trainloader_l_iter
# Check if the loader has a batch available
try:
img_u,mask_u,ohmask_u,_ = next(loader_u)
except:
trainloader_u_iter = iter(trainloader_u)
loader_u = trainloader_u_iter
img_u,mask_u,ohmask_u,_ = next(loader_u)
if args.nogpu:
img_u,mask_u,ohmask_u = Variable(img_u),Variable(mask_u),Variable(ohmask_u)
else:
img_u,mask_u,ohmask_u = Variable(img_u.cuda()),Variable(mask_u.cuda()),Variable(ohmask_u.cuda())
# semi unlabelled training
cpmap = generator(img_u)
cpmapsmax = nn.Softmax2d()(cpmap)
conf = discriminator(cpmapsmax)
confsmax = nn.Softmax2d()(conf)
conflsmax = nn.LogSoftmax()(conf)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Adversarial Loss
targetr = Variable(torch.ones((N,H,W)).long())
if not args.nogpu:
targetr = targetr.cuda()
LGadv = nn.NLLLoss2d()(conflsmax,targetr)
LGadv_d = LGadv.data
hardpred = torch.max(cpmapsmax,1)[1].squeeze(1)
idx = np.zeros(cpmap.data.cpu().numpy().shape,dtype=np.uint8)
idx = idx.transpose(0, 2, 3, 1)
confnp = confsmax[:,1,...].data.cpu().numpy()
hardprednp = hardpred.data.cpu().numpy()
idx[confnp > args.t_semi] = np.identity(2, dtype=idx.dtype)[hardprednp[ confnp > args.t_semi]]
LG = args.lam_adv*LGadv
if np.count_nonzero(idx) != 0:
cpmaplsmax = nn.LogSoftmax()(cpmap)
idx = Variable(torch.from_numpy(idx.transpose(0,3,1,2)).byte().cuda())
LGsemi_arr = cpmaplsmax.masked_select(idx.bool())
LGsemi = -1*LGsemi_arr.mean()
LGsemi_d = LGsemi.data
# LGsemi_c = LGsemi.cpu().numpy()
LGsemi = args.lam_semi*LGsemi
# LG += LGsemi
LG.backward()
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
################################################
# train labelled data #
################################################
loader_l = trainloader_l_iter
loader_u = trainloader_u_iter
try:
if random.random() <0.5:
img_l,mask_l,ohmask_l,_ = next(loader_l)
else:
img_l,mask_l,ohmask_l,_ = next(loader_u)
except:
break
if args.nogpu:
img_l,mask_l,ohmask_l = Variable(img_l),Variable(mask_l),Variable(ohmask_l)
else:
img_l,mask_l,ohmask_l = Variable(img_l.cuda()),Variable(mask_l.cuda()),Variable(ohmask_l.cuda())
################################################
# Labelled data for Discriminator Training #
################################################
cpmap = generator(Variable(img_l.data,volatile=True))
cpmap = nn.Softmax2d()(cpmap)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N,H,W)).long())
targetr = Variable(torch.ones((N,H,W)).long())
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmask_l.float()))
optimD.zero_grad()
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth)*nn.NLLLoss2d()(confr,targetr) # mask_l
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr,targetf) # targetf mask_l
else:
LDr = nn.NLLLoss2d()(confr, targetr) # targetr mask_l
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data)))
LDf = nn.NLLLoss2d()(conff,mask_l) # targetf
LDf.backward()
LDr_d = LDr.data
# LDr_c = LDr.cpu().numpy()
LDf_d = LDf.data
# LDf_c = LDf.cpu().numpy()
LD_d = LDr_d + LDf_d
# LD_c = LDr_c + LDf_c
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
#####################################
# labelled data Generator Training #
#####################################
optimG.zero_grad()
optimD.zero_grad()
cpmap = generator(img_l)
# print("cpmap: ", cpmap.size())
cpmapsmax = nn.Softmax2d()(cpmap)
# print("cpmapsmax: ", cpmapsmax.size())
cpmaplsmax = nn.LogSoftmax()(cpmap)
# print("cpmaplsmax: ", cpmaplsmax.size())
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax,mask_l)
LGadv = nn.NLLLoss2d()(conff,targetr) # targetr mask_l
LGadv_d = LGadv.data
# LGadv_c = LGadv.cpu().numpy()
LGce_d = LGce.data
# LGce_c = LGce.cpu().numpy()
# LGsemi_d = 0 # No semi-supervised training
LGadv = args.lam_adv*LGadv
(LGce + LGadv).backward()
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
LGseg_d = LGce_d + LGadv_d + LGsemi_d
# LGseg_c = LGce_c + LGadv_c + LGsemi_c
# training_log = [epoch,itr,LD_c,LDr_c,LDf_c,LGseg_c,LGce_c,LGadv_c,LGsemi_c]
print("[{}][{}] LD: {:.4f} LD_fake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f} LG_semi: {:.4f}"\
.format(epoch,itr,LD_d,LDr_d,LDf_d,LGseg_d,LGce_d,LGadv_d,LGsemi_d))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou, best_eiou = snapshot_segdis(generator,discriminator,valoader,epoch,best_miou,best_eiou,args.snapshot_dir,args.prefix)
def train_adv(generator,discriminator,optimG,optimD,trainloader,valoader,args,ws):
best_eiou = -1
best_miou = -1
for epoch in range(args.start_epoch,args.max_epoch+1):
generator.train()
for batch_id, (img,mask,ohmask,_) in enumerate(trainloader):
if args.nogpu:
img,mask,ohmask = Variable(img),Variable(mask),Variable(ohmask)
else:
img,mask,ohmask = Variable(img.cuda()),Variable(mask.cuda()),Variable(ohmask.cuda())
itr = len(trainloader)*(epoch-1) + batch_id
# generator forward
cpmap = generator(Variable(img.data,volatile=True))
cpmap = nn.Softmax2d()(cpmap)
# print("cpmap: ", cpmap.size(), " ohmask: ", ohmask.size())
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# print("cpmap: ", cpmap.size())
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N,H,W)).long(),requires_grad=False)
targetr = Variable(torch.ones((N,H,W)).long(),requires_grad=False)
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
##########################
# DISCRIMINATOR TRAINING #
##########################
optimD.zero_grad()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmask.float()))
# print("confr: ", confr.size())
# print("targetr: ", targetr.size())
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth)*nn.NLLLoss2d()(confr,targetr)
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr,targetf)
else:
LDr = nn.NLLLoss2d()(confr,targetr)
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data)))
LDf = nn.NLLLoss2d()(conff,targetf)
LDf.backward()
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
######################
# GENERATOR TRAINING #
#####################
optimG.zero_grad()
optimD.zero_grad()
cmap = generator(img)
cpmapsmax = nn.Softmax2d()(cmap)
cpmaplsmax = nn.LogSoftmax()(cmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax,mask)
LGadv = nn.NLLLoss2d()(conff,targetr)
LGseg = LGce + args.lam_adv *LGadv
LGseg.backward()
poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
print("[{}][{}] LD: {:.4f} LDfake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f}" \
.format(epoch,itr,(LDr + LDf).data,LDr.data,LDf.data,LGseg.data,LGce.data,LGadv.data))
# .format(epoch,itr,(LDr + LDf).data[0],LDr.data[0],LDf.data[0],LGseg.data[0],LGce.data[0],LGadv.data[0]))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou, best_eiou = snapshot_segdis(generator,discriminator,valoader,epoch,best_miou,best_eiou,args.snapshot_dir,args.prefix)
def train_semi_m(generator, discriminator, optimG, optimD, trainloader_l,
trainloader_u, valoader, args):
best_miou = -1
# bce_loss = BCEWithLogitsLoss2d()
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
if epoch > args.wait_semi:
"""
Using Discriminator ouput to train generator
"""
loss_semi_value = 0
for batch_id, (img, mask, ohmask) in enumerate(trainloader_u):
if args.nogpu:
img, mask, ohmask = Variable(img), Variable(
mask), Variable(ohmask)
else:
img, mask, ohmask = Variable(img.cuda()), Variable(
mask.cuda()), Variable(ohmask.cuda())
# generator output
cpmap = generator(Variable(img.data, volatile=True))
# cpmap.detach()
cpmapsmax = nn.Softmax2d()(
cpmap) # torch.Size([4, 2, 320, 320])
# discriminator output
conf = discriminator(cpmapsmax)
# confsmax = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1)
confsmax = nn.Softmax2d()(conf)
conflsmax = nn.LogSoftmax()(
conf) # torch.Size([4, 2, 320, 320])
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Adversarial Loss
targetr = Variable(torch.ones((N, H, W)).long())
if not args.nogpu:
targetr = targetr.cuda()
LGadv = nn.NLLLoss2d()(conflsmax, targetr)
LGadv_d = LGadv.data
# D_output = confsmax.data.cpu().numpy().squeeze(axis=1)
D_output = torch.max(confsmax, 1)[1].cpu().numpy()
# print("D_output: ", D_output.shape)
semi_ignore_mask = (D_output < args.t_semi)
semi_gt = cpmapsmax.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255 # 255
semi_ratio = 1.0 - float(
semi_ignore_mask.sum()) / semi_ignore_mask.size
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
LG_semi = args.lam_semi * loss_calc(cpmapsmax, semi_gt)
LG_semi = LG_semi # /args.iter_size
# loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi
loss_semi_value += LG_semi.data.cpu().numpy(
) / args.lam_semi
# LG_semi += LGadv_d
LG_semi.backward()
# LGsemi.backward()
# optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
# optimG.step()
for batch_id, (img, mask, ohmask) in enumerate(trainloader_l):
if args.nogpu:
img, mask, ohmask = Variable(img), Variable(mask), Variable(
ohmask)
else:
img, mask, ohmask = Variable(img.cuda()), Variable(
mask.cuda()), Variable(ohmask.cuda())
# itr = len(trainloader)*(epoch-1) + batch_id
itr = (len(trainloader_u) + len(trainloader_l)) * (epoch -
1) + batch_id
# generator forward
cpmap = generator(Variable(img.data, volatile=True))
cpmap = nn.Softmax2d()(cpmap)
# print("cpmap: ", cpmap.size(), " ohmask: ", ohmask.size())
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# print("cpmap: ", cpmap.size())
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N, H, W)).long(),
requires_grad=False)
targetr = Variable(torch.ones((N, H, W)).long(),
requires_grad=False)
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
##########################
# DISCRIMINATOR TRAINING #
##########################
optimD.zero_grad()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmask.float()))
# print("confr: ", confr.size())
# print("targetr: ", targetr.size())
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(confr,
targetr)
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr, targetf)
else:
LDr = nn.NLLLoss2d()(confr, targetr)
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data)))
LDf = nn.NLLLoss2d()(conff, targetf)
LDf.backward()
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
######################
# GENERATOR TRAINING #
#####################
optimG.zero_grad()
optimD.zero_grad()
cmap = generator(img)
cpmapsmax = nn.Softmax2d()(cmap)
cpmaplsmax = nn.LogSoftmax()(cmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax, mask)
LGadv = nn.NLLLoss2d()(conff, targetr)
LGseg = LGce + args.lam_adv * LGadv
LGseg.backward()
poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
print("[{}][{}] LD: {:.4f} LDfake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f}, LG_semi: {:.5f}" \
.format(epoch,itr,(LDr + LDf).data,LDr.data,LDf.data,LGseg.data,LGce.data,LGadv.data,LG_semi.data))
# .format(epoch,itr,(LDr + LDf).data[0],LDr.data[0],LDf.data[0],LGseg.data[0],LGce.data[0],LGadv.data[0]))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou = snapshot_segdis(generator, discriminator, valoader, epoch,
best_miou, args.snapshot_dir, args.prefix)
def train_semi(generator, discriminator, optimG, optimD, trainloader_l,
trainloader_u, valoader, args):
best_miou = -1
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
trainloader_l_iter = iter(trainloader_l)
trainloader_u_iter = iter(trainloader_u)
print("Epoch: {}".format(epoch))
batch_id = 0
# Randomly pick labeled or unlabeled data for training
while (True):
if random.random() < 0.5:
loader = trainloader_l_iter
labeled = True
else:
loader = trainloader_u_iter
labeled = False
# Check if the loader has a batch available
try:
img, mask, ohmask = next(loader)
except:
# Curr loader doesn't have data
if labeled:
loader = trainloader_u_iter
labeled = False
else:
loader = trainloader_l_iter
labeled = True
# Check if the new loader has data
try:
img, mask, ohmask = next(loader)
except:
# Boith loaders exhausted
break
batch_id += 1
if args.nogpu:
img, mask, ohmask = Variable(img), Variable(mask), Variable(
ohmask)
else:
img, mask, ohmask = Variable(img.cuda()), Variable(
mask.cuda()), Variable(ohmask.cuda())
itr = (len(trainloader_u) + len(trainloader_l)) * (epoch -
1) + batch_id
LGseg_d = 0
if epoch < args.wait_semi:
################################################
# Labelled data for Discriminator Training #
################################################
cpmap = generator(Variable(img.data, volatile=True))
cpmap = nn.Softmax2d()(cpmap)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N, H, W)).long())
targetr = Variable(torch.ones((N, H, W)).long())
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmask.float()))
optimD.zero_grad()
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(confr,
targetr)
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr, targetf)
else:
LDr = nn.NLLLoss2d()(confr, targetr)
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data)))
LDf = nn.NLLLoss2d()(conff, targetf)
LDf.backward()
LDr_d = LDr.data
LDf_d = LDf.data
LD_d = LDr_d + LDf_d
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
#####################################
# labelled data Generator Training #
#####################################
optimG.zero_grad()
optimD.zero_grad()
cpmap = generator(img)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax, mask)
LGadv = nn.NLLLoss2d()(conff, targetr)
LGadv_d = LGadv.data
LGce_d = LGce.data
LGsemi_d = 0 # No semi-supervised training
LGadv = args.lam_adv * LGadv
(LGce + LGadv).backward()
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
LGseg_d = LGce_d + args.lam_adv * LGadv_d + args.lam_semi * LGsemi_d
else:
if labeled:
################################################
# Labelled data for Discriminator Training #
################################################
cpmap = generator(Variable(img.data, volatile=True))
cpmap = nn.Softmax2d()(cpmap)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N, H, W)).long())
targetr = Variable(torch.ones((N, H, W)).long())
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmask.float()))
optimD.zero_grad()
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(
confr, targetr)
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr,
targetf)
else:
LDr = nn.NLLLoss2d()(confr, targetr)
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(
cpmap.data)))
LDf = nn.NLLLoss2d()(conff, targetf)
LDf.backward()
LDr_d = LDr.data
LDf_d = LDf.data
LD_d = LDr_d + LDf_d
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
#####################################
# labelled data Generator Training #
#####################################
optimG.zero_grad()
optimD.zero_grad()
cpmap = generator(img)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax, mask)
LGadv = nn.NLLLoss2d()(conff, targetr)
LGadv_d = LGadv.data
LGce_d = LGce.data
LGsemi_d = 0 # No semi-supervised training
LGadv = args.lam_adv * LGadv
(LGce + LGadv).backward()
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
LGseg_d = LGce_d + args.lam_adv * LGadv_d + args.lam_semi * LGsemi_d
else:
#####################################
# Use unlabelled data to get L_semi #
#####################################
# No discriminator training
LD_d = 0
LDr_d = 0
LDf_d = 0
# Init all loss to 0 for logging ease
LGsemi_d = 0
LGce_d = 0
LGadv_d = 0
optimG.zero_grad()
# if epoch > args.wait_semi:
cpmap = generator(img)
cpmapsmax = nn.Softmax2d()(cpmap)
conf = discriminator(cpmapsmax)
confsmax = nn.Softmax2d()(conf)
conflsmax = nn.LogSoftmax()(conf)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Adversarial Loss
targetr = Variable(torch.ones((N, H, W)).long())
if not args.nogpu:
targetr = targetr.cuda()
LGadv = nn.NLLLoss2d()(conflsmax, targetr)
LGadv_d = LGadv.data
# Semi-Supervised Loss
hardpred = torch.max(cpmapsmax, 1)[1].squeeze(1)
idx = np.zeros(cpmap.data.cpu().numpy().shape,
dtype=np.uint8)
idx = idx.transpose(0, 2, 3, 1)
confnp = confsmax[:, 1, ...].data.cpu().numpy()
hardprednp = hardpred.data.cpu().numpy()
idx[confnp > args.t_semi] = np.identity(
2, dtype=idx.dtype)[hardprednp[confnp > args.t_semi]]
LG = args.lam_adv * LGadv
if np.count_nonzero(idx) != 0:
cpmaplsmax = nn.LogSoftmax()(cpmap)
idx = Variable(
torch.from_numpy(idx.transpose(0, 3, 1,
2)).byte().cuda())
# LGsemi_arr = cpmaplsmax.masked_select(idx)
LGsemi_arr = cpmaplsmax.masked_select(idx.bool())
LGsemi = -1 * LGsemi_arr.mean()
LGsemi_d = LGsemi.data
LGsemi = args.lam_semi * LGsemi
LG += LGsemi
LG.backward()
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
# Manually free all variables. Look into details of how variables are freed
del idx
del confnp
del confsmax
del hardpred
del hardprednp
del cpmapsmax
del cpmap
LGseg_d = LGce_d + args.lam_adv * LGadv_d + args.lam_semi * LGsemi_d
# Manually free memory! Later, really understand how computation graphs free variables
# LD: Discriminator loss LD_fake: fake Dis los LD_real: real Dis loss LG_ce: generator loss
# Ladv: Adversarail loss Lsemi: semi loss
print("[{}][{}] LD: {:.4f} LD_fake: {:.4f} LD_real: {:.4f} LGseg: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f} LG_semi: {:.4f}"\
.format(epoch,itr,LD_d,LDf_d,LDr_d,LGseg_d,LGce_d,LGadv_d,LGsemi_d))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou = snapshot_segdis(generator, discriminator, valoader, epoch,
best_miou, args.snapshot_dir, args.prefix)
def train_label_correction(generator, optimG, trainloader, valoader, args):
best_miou = -1
best_eiou = -1
# FILE_INFO = file_info()
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
for batch_id, (img_l, mask_l, ohmask_l, _,
imgs_org) in enumerate(trainloader):
itr = len(trainloader) * (epoch - 1) + batch_id
if args.nogpu:
img_l, mask_l, ohmask_l = Variable(img_l), Variable(
mask_l), Variable(ohmask_l)
else:
img_l, mask_l, ohmask_l = Variable(img_l.cuda()), Variable(
mask_l.cuda()), Variable(ohmask_l.cuda())
imgs_array = imgs_org.numpy()
LGcorr = torch.tensor(0)
loss_corr_value = 0
# begin label correction
################################################
# train labelled data #
################################################
cpmap = generator(img_l)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
LGce = nn.NLLLoss()(cpmaplsmax, mask_l)
LGce_d = LGce.data
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.zero_grad()
LGce.backward()
optimG.step()
################################################
# label correction #
################################################
if epoch > args.wait_semi:
# Check if the loader has a batch available
# semi unlabelled training
cpmap = generator(img_l)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
soft_preds_np = cpmapsmax.data.cpu().numpy()
soft_mask = np.zeros(
(soft_preds_np.shape[0], soft_preds_np.shape[2],
soft_preds_np.shape[3]))
for i in range(soft_preds_np.shape[0]):
soft_pred = soft_preds_np[i]
img_array = imgs_array[i]
img_array = cv2.resize(img_array, (320, 320),
interpolation=cv2.INTER_AREA)
hard_pred = np.argmax(soft_pred, axis=0).astype(np.uint8)
rois_crf, rois = get_roi_crf(img_array, soft_pred,
hard_pred)
for roi_crf, roi in zip(rois_crf, rois):
# print(roi)
xmin, ymin, xmax, ymax = roi[0], roi[1], roi[2], roi[3]
soft_pred[:, ymin:ymax, xmin:xmax] = roi_crf
crf_result = np.argmax(soft_pred, axis=0).astype(np.uint8)
soft_mask[i] = crf_result
mask_c = Variable(torch.tensor(soft_mask)).type(
torch.LongTensor).cuda()
LGcorr = nn.NLLLoss()(cpmaplsmax, mask_c)
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.zero_grad()
LGcorr.backward()
optimG.step()
LGcorr_d = LGcorr.data
LGseg_d = LGce_d + args.lam_adv * LGcorr
print("[{}][{}] LGseg_d: {:.4f} LGce_d: {:.4f} LG_corr: {:.4f}"\
.format(epoch,itr,LGseg_d,LGce_d,LGcorr_d))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou, best_eiou = snapshote(generator, valoader, epoch, best_miou,
best_eiou, args.snapshot_dir,
args.prefix)
def train_semi(generator, discriminator, optimG, optimD, trainloader_l,
trainloader_u, valoader, args):
best_miou = -1
best_eiou = -1
FILE_INFO = file_info()
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
trainloader_l_iter = iter(trainloader_l)
trainloader_u_iter = iter(trainloader_u)
print("Epoch: {}".format(epoch))
batch_id = 0
while (True):
batch_id += 1
itr = (len(trainloader_u) + len(trainloader_l)) * (epoch -
1) + batch_id
optimG.zero_grad()
optimD.zero_grad()
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
loader_l = trainloader_l_iter
try:
img_l, mask_l, ohmask_l, _ = next(loader_l)
except:
break
if args.nogpu:
img_l, mask_l, ohmask_l = Variable(img_l), Variable(
mask_l), Variable(ohmask_l)
else:
img_l, mask_l, ohmask_l = Variable(img_l.cuda()), Variable(
mask_l.cuda()), Variable(ohmask_l.cuda())
LGsemi_d = 0
loss_semi_value = 0
if epoch > args.wait_semi:
# Check if the loader has a batch available
# semi unlabelled training
cpmap = generator(img_l)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
# confsmax = nn.Softmax2d()(conf)
conf = discriminator(cpmapsmax)
conf_softmax = nn.Softmax2d()(conf)
conf_softmax_remain = conf_softmax.detach().squeeze(1)
conf_soft = conf_softmax_remain.data.cpu().numpy()
semi_gt = np.argmax(conf_soft, axis=1).astype(np.uint8)
semi_ratio = 1.0 - float(semi_gt.sum()) / semi_gt.size
# print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.LongTensor(semi_gt).cuda()
LGsemi = args.lam_semi * nn.NLLLoss()(cpmaplsmax, semi_gt)
LGsemi_d += LGsemi.data.cpu().numpy() / args.lam_semi
# LGsemi += LGadv_d
LGsemi.backward()
################################################
# train labelled data #
################################################
cpmap = generator(img_l)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss()(cpmaplsmax, mask_l)
# LGce = loss_calc(cpmaplsmax, mask_l)
# LGadv = nn.BCELoss()(conff,mask_l)
# LGadv = loss_calc(conff, mask_l)
LGadv = nn.NLLLoss()(conff, mask_l)
LGadv_d = LGadv.data
LGce_d = LGce.data
# print("{}: {}".format(str(FILE_INFO), conff.size()))
# LGsemi_d = 0 # No semi-supervised training
LGadv = args.lam_adv * LGadv
(LGce + LGadv).backward()
# optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
optimD.step()
LGseg_d = LGce_d + LGadv_d + LGsemi_d
# LGseg_c = LGce_c + LGadv_c + LGsemi_c
# training_log = [epoch,itr,LD_c,LDr_c,LDf_c,LGseg_c,LGce_c,LGadv_c,LGsemi_c]
print("[{}][{}] LGseg_d: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f} LG_semi: {:.4f}"\
.format(epoch,itr,LGseg_d,LGce_d,LGadv_d,LGsemi_d))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou, best_eiou = snapshot_segdis(generator, discriminator,
valoader, epoch, best_miou,
best_eiou, args.snapshot_dir,
args.prefix)
def train_semi_b(generator, discriminator, optimG, optimD, trainloader_l,
trainloader_u, valoader, args):
best_miou = -1
best_eiou = -1
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
trainloader_l_iter = iter(trainloader_l)
trainloader_u_iter = iter(trainloader_u)
print("Epoch: {}".format(epoch))
batch_id = 0
while (True):
batch_id += 1
itr = (len(trainloader_u) + len(trainloader_l)) * (epoch -
1) + batch_id
optimG.zero_grad()
optimD.zero_grad()
optimD = poly_lr_scheduler(optimD, args.d_lr, itr)
optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
LGsemi_d = 0
if epoch > args.wait_semi:
if random.random() < 0.5:
loader_l = trainloader_l_iter
loader_u = trainloader_u_iter
else:
loader_l = trainloader_u_iter
loader_u = trainloader_l_iter
# Check if the loader has a batch available
try:
img_u, mask_u, ohmask_u, _ = next(loader_u)
except:
trainloader_u_iter = iter(trainloader_u)
loader_u = trainloader_u_iter
img_u, mask_u, ohmask_u, _ = next(loader_u)
if args.nogpu:
img_u, mask_u, ohmask_u = Variable(img_u), Variable(
mask_u), Variable(ohmask_u)
else:
img_u, mask_u, ohmask_u = Variable(img_u.cuda()), Variable(
mask_u.cuda()), Variable(ohmask_u.cuda())
# semi unlabelled training
cpmap = generator(img_u)
cpmapsmax = nn.Softmax2d()(cpmap)
# confsmax = nn.Softmax2d()(conf)
conf = discriminator(cpmapsmax)
confsigmoid = nn.Sigmoid()(conf)
confsigmoid_remain = confsigmoid.detach().squeeze(1)
# conflsmax = nn.LogSoftmax()(conf)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
semi_gt = np.int8(
(confsigmoid_remain.cpu().numpy() >= args.t_semi))
# semi_gt = cpmap.data.cpu().numpy().argmax(axis=1)
# semi_gt[semi_ignore_mask] = 1
semi_ratio = 1.0 - float(semi_gt.sum()) / semi_gt.size
# print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
LGsemi = args.lam_semi * loss_calc(cpmapsmax, semi_gt)
LGsemi_d += LGsemi.data.cpu().numpy() / args.lam_semi
# LGsemi += LGadv_d
LGsemi.backward()
################################################
# train labelled data #
################################################
loader_l = trainloader_l_iter
loader_u = trainloader_u_iter
try:
if random.random() < 0.5:
img_l, mask_l, ohmask_l, _ = next(loader_l)
else:
img_l, mask_l, ohmask_l, _ = next(loader_u)
except:
break
if args.nogpu:
img_l, mask_l, ohmask_l = Variable(img_l), Variable(
mask_l), Variable(ohmask_l)
else:
img_l, mask_l, ohmask_l = Variable(img_l.cuda()), Variable(
mask_l.cuda()), Variable(ohmask_l.cuda())
#####################################
# labelled data Generator Training #
#####################################
# optimG.zero_grad()
# optimD.zero_grad()
cpmap = generator(img_l)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N, H, W))) # .long()
targetr = Variable(torch.ones((N, H, W))) # .long()
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
conff = nn.Sigmoid()(discriminator(cpmapsmax))
print("cpmaplsmax: ", cpmaplsmax.size())
print("mask_l: ", mask_l.size())
print("conff: ", conff.size())
# LGce = loss_calc(cpmaplsmax, mask_l)
LGce = nn.NLLLoss()(cpmaplsmax, mask_l)
# LGadv = nn.BCELoss()(conff,mask_l)
LGadv = loss_calc(cpmaplsmax, mask_l)
LGadv_d = LGadv.data
# LGadv_c = LGadv.cpu().numpy()
LGce_d = LGce.data
# LGce_c = LGce.cpu().numpy()
# LGsemi_d = 0 # No semi-supervised training
LGadv = args.lam_adv * LGadv
(LGce + LGadv).backward()
# optimG = poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
optimD.step()
LGseg_d = LGce_d + LGadv_d + LGsemi_d
# LGseg_c = LGce_c + LGadv_c + LGsemi_c
# training_log = [epoch,itr,LD_c,LDr_c,LDf_c,LGseg_c,LGce_c,LGadv_c,LGsemi_c]
print("[{}][{}] LGseg_d: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f} LG_semi: {:.4f}"\
.format(epoch,itr,LGseg_d,LGce_d,LGadv_d,LGsemi_d))
# best_miou = snapshot(generator,valoader,epoch,best_miou,args.snapshot_dir,args.prefix)
best_miou, best_eiou = snapshot_segdis(generator, discriminator,
valoader, epoch, best_miou,
best_eiou, args.snapshot_dir,
args.prefix)
def train_semi(args):
# TODO: Make it more generic to include for other splits
args.batch_size = args.batch_size * 2
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop((321, 321))]
trainset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), label_transform=Compose(labtr), \
co_transform=Compose(cotr))
trainloader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
drop_last=True)
#########################
# Validation Dataloader #
########################
if args.val_orig:
if args.no_norm:
imgtr = [ZeroPadding(), ToTensor()]
else:
imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = []
else:
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop((321, 321))]
valset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \
label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False)
valoader = DataLoader(valset, batch_size=1)
#############
# GENERATOR #
#############
generator = deeplabv2.ResDeeplab()
optimG = optim.SGD(filter(lambda p: p.requires_grad, \
generator.parameters()),lr=args.g_lr,momentum=0.9,\
weight_decay=0.0001,nesterov=True)
if not args.nogpu:
generator = nn.DataParallel(generator).cuda()
#################
# DISCRIMINATOR #
################
discriminator = Dis(in_channels=21)
if args.d_optim == 'adam':
optimD = optim.Adam(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr = args.d_lr)
else:
optimD = optim.SGD(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.5,nesterov=True)
if not args.nogpu:
discriminator = nn.DataParallel(discriminator).cuda()
############
# TRAINING #
############
best_miou = -1
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
for batch_id, (img, mask, ohmask) in enumerate(trainloader):
if args.nogpu:
img, mask, ohmask = Variable(img), Variable(mask), Variable(
ohmask)
else:
img, mask, ohmask = Variable(img.cuda()), Variable(
mask.cuda()), Variable(ohmask.cuda())
itr = len(trainloader) * (epoch - 1) + batch_id
## TODO: Extend random interleaving for split of any size
mid = args.batch_size // 2
img1, mask1, ohmask1 = img[0:mid, ...], mask[0:mid,
...], ohmask[0:mid,
...]
img2, mask2, ohmask2 = img[mid:, ...], mask[mid:,
...], ohmask[mid:, ...]
# Random Interleaving
if random.random() < 0.5:
imgl, maskl, ohmaskl = img1, mask1, ohmask1
imgu, masku, ohmasku = img2, mask2, ohmask2
else:
imgu, masku, ohmasku = img1, mask1, ohmask1
imgl, maskl, ohmaskl = img2, mask2, ohmask2
################################################
# Labelled data for Discriminator Training #
################################################
cpmap = generator(Variable(imgl.data, volatile=True))
cpmap = nn.Softmax2d()(cpmap)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N, H, W)).long())
targetr = Variable(torch.ones((N, H, W)).long())
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmaskl.float()))
optimD.zero_grad()
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(confr,
targetr)
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr, targetf)
else:
LDr = nn.NLLLoss2d()(confr, targetr)
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data)))
LDf = nn.NLLLoss2d()(conff, targetf)
LDf.backward()
poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
###########################################
# labelled data Generator Training #
###########################################
optimG.zero_grad()
cpmap = generator(imgl)
cpmapsmax = nn.Softmax2d()(cpmap)
cpmaplsmax = nn.LogSoftmax()(cpmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax, maskl)
LGadv = nn.NLLLoss2d()(conff, targetr)
LGadv_d = LGadv.data[0]
LGce_d = LGce.data[0]
LGadv = args.lam_adv * LGadv
(LGce + LGadv).backward()
#####################################
# Use unlabelled data to get L_semi #
#####################################
LGsemi_d = 0
if epoch > args.wait_semi:
cpmap = generator(imgu)
softpred = nn.Softmax2d()(cpmap)
hardpred = torch.max(softpred, 1)[1].squeeze(1)
conf = nn.Softmax2d()(discriminator(
Variable(softpred.data, volatile=True)))
idx = np.zeros(cpmap.data.cpu().numpy().shape, dtype=np.uint8)
idx = idx.transpose(0, 2, 3, 1)
confnp = cpmap[:, 1, ...].data.cpu().numpy()
hardprednp = hardpred.data.cpu().numpy()
idx[confnp > args.t_semi] = np.identity(
21, dtype=idx.dtype)[hardprednp[confnp > args.t_semi]]
if np.count_nonzero(idx) != 0:
cpmaplsmax = nn.LogSoftmax()(cpmap)
idx = Variable(torch.from_numpy(idx).byte().cuda())
LGsemi_arr = cpmaplsmax.masked_select(idx)
LGsemi = -1 * LGsemi_arr.mean()
LGsemi_d = LGsemi.data[0]
LGsemi = args.lam_semi * LGsemi
LGsemi.backward()
else:
LGsemi_d = 0
LGseg_d = LGce_d + LGadv_d + LGsemi_d
del idx
del conf
del confnp
del hardpred
del softpred
del hardprednp
del cpmap
LGseg_d = LGce_d + LGadv_d + LGsemi_d
poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
# Manually free memory! Later, really understand how computation graphs free variables
print("[{}][{}] LD: {:.4f} LD_fake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f} LG_semi: {:.4f}"\
.format(epoch,itr,(LDr + LDf).data[0],LDr.data[0],LDf.data[0],LGseg_d,LGce_d,LGadv_d,LGsemi_d))
snapshot(generator, valoader, epoch, best_miou, args.snapshot_dir,
args.prefix)
def train_adv(args):
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop((321, 321))]
trainset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), label_transform=Compose(labtr), \
co_transform=Compose(cotr))
trainloader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
drop_last=True)
#########################
# Validation Dataloader #
########################
if args.val_orig:
if args.no_norm:
imgtr = [ZeroPadding(), ToTensor()]
else:
imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = []
else:
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop((321, 321))]
valset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \
label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False)
valoader = DataLoader(valset, batch_size=1)
#############
# GENERATOR #
#############
generator = deeplabv2.ResDeeplab()
optimG = optim.SGD(filter(lambda p: p.requires_grad, \
generator.parameters()),lr=args.g_lr,momentum=0.9,\
weight_decay=0.0001,nesterov=True)
if not args.nogpu:
generator = nn.DataParallel(generator).cuda()
#################
# DISCRIMINATOR #
################
discriminator = Dis(in_channels=21)
if args.d_optim == 'adam':
optimD = optim.Adam(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr = args.d_lr)
else:
optimD = optim.SGD(filter(lambda p: p.requires_grad, \
discriminator.parameters()),lr=args.d_lr,weight_decay=0.0001,momentum=0.5,nesterov=True)
if not args.nogpu:
discriminator = nn.DataParallel(discriminator).cuda()
#############
# TRAINING #
#############
best_miou = -1
for epoch in range(args.start_epoch, args.max_epoch + 1):
generator.train()
for batch_id, (img, mask, ohmask) in enumerate(trainloader):
if args.nogpu:
img, mask, ohmask = Variable(img), Variable(mask), Variable(
ohmask)
else:
img, mask, ohmask = Variable(img.cuda()), Variable(
mask.cuda()), Variable(ohmask.cuda())
itr = len(trainloader) * (epoch - 1) + batch_id
cpmap = generator(Variable(img.data, volatile=True))
cpmap = nn.Softmax2d()(cpmap)
N = cpmap.size()[0]
H = cpmap.size()[2]
W = cpmap.size()[3]
# Generate the Real and Fake Labels
targetf = Variable(torch.zeros((N, H, W)).long(),
requires_grad=False)
targetr = Variable(torch.ones((N, H, W)).long(),
requires_grad=False)
if not args.nogpu:
targetf = targetf.cuda()
targetr = targetr.cuda()
##########################
# DISCRIMINATOR TRAINING #
##########################
optimD.zero_grad()
# Train on Real
confr = nn.LogSoftmax()(discriminator(ohmask.float()))
if args.d_label_smooth != 0:
LDr = (1 - args.d_label_smooth) * nn.NLLLoss2d()(confr,
targetr)
LDr += args.d_label_smooth * nn.NLLLoss2d()(confr, targetf)
else:
LDr = nn.NLLLoss2d()(confr, targetr)
LDr.backward()
# Train on Fake
conff = nn.LogSoftmax()(discriminator(Variable(cpmap.data)))
LDf = nn.NLLLoss2d()(conff, targetf)
LDf.backward()
poly_lr_scheduler(optimD, args.d_lr, itr)
optimD.step()
######################
# GENERATOR TRAINING #
#####################
optimG.zero_grad()
cmap = generator(img)
cpmapsmax = nn.Softmax2d()(cmap)
cpmaplsmax = nn.LogSoftmax()(cmap)
conff = nn.LogSoftmax()(discriminator(cpmapsmax))
LGce = nn.NLLLoss2d()(cpmaplsmax, mask)
LGadv = nn.NLLLoss2d()(conff, targetr)
LGseg = LGce + args.lam_adv * LGadv
LGseg.backward()
poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.step()
print("[{}][{}] LD: {:.4f} LDfake: {:.4f} LD_real: {:.4f} LG: {:.4f} LG_ce: {:.4f} LG_adv: {:.4f}" \
.format(epoch,itr,(LDr + LDf).data[0],LDr.data[0],LDf.data[0],LGseg.data[0],LGce.data[0],LGadv.data[0]))
snapshot(generator, valoader, epoch, best_miou, args.snapshot_dir,
args.prefix)
def train_base(args):
#######################
# Training Dataloader #
#######################
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop((321, 321))]
trainset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), label_transform=Compose(labtr), \
co_transform=Compose(cotr))
trainloader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2,
drop_last=True)
#########################
# Validation Dataloader #
########################
if args.val_orig:
if args.no_norm:
imgtr = [ZeroPadding(), ToTensor()]
else:
imgtr = [ZeroPadding(), ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = []
else:
if args.no_norm:
imgtr = [ToTensor()]
else:
imgtr = [ToTensor(), NormalizeOwn()]
labtr = [IgnoreLabelClass(), ToTensorLabel()]
cotr = [RandomSizedCrop((321, 321))]
valset = PascalVOC(home_dir,args.dataset_dir,img_transform=Compose(imgtr), \
label_transform = Compose(labtr),co_transform=Compose(cotr),train_phase=False)
valoader = DataLoader(valset, batch_size=1)
model = deeplabv2.ResDeeplab()
init_weights(model, args.init_net)
optimG = optim.SGD(filter(lambda p: p.requires_grad, \
model.parameters()),lr=args.g_lr,momentum=0.9,\
weight_decay=0.0001,nesterov=True)
if not args.nogpu:
model = nn.DataParallel(model).cuda()
best_miou = -1
for epoch in range(args.start_epoch, args.max_epoch + 1):
model.train()
for batch_id, (img, mask, _) in enumerate(trainloader):
if args.nogpu:
img, mask = Variable(img), Variable(mask)
else:
img, mask = Variable(img.cuda()), Variable(mask.cuda())
itr = len(trainloader) * (epoch - 1) + batch_id
cprob = model(img)
cprob = nn.LogSoftmax()(cprob)
Lseg = nn.NLLLoss2d()(cprob, mask)
poly_lr_scheduler(optimG, args.g_lr, itr)
optimG.zero_grad()
Lseg.backward()
optimG.step()
print("[{}][{}]Loss: {:0.4f}".format(epoch, itr, Lseg.data[0]))
snapshot(model, valoader, epoch, best_miou, args.snapshot_dir,
args.prefix)
