# torch.manual_seed(1234)

# torch.cuda.manual_seed(1234)

opt = TrainOptions()

args = opt.initialize()

_t = {'iter time' : Timer()}

model_name = args.source + '_to_' + args.target

if not os.path.exists(args.snapshot_dir):

os.makedirs(args.snapshot_dir)

os.makedirs(os.path.join(args.snapshot_dir, 'logs'))

opt.print_options(args)

sourceloader, targetloader = CreateSrcDataLoader(args), CreateTrgDataLoader(args)

targetloader_iter, sourceloader_iter = iter(targetloader), iter(sourceloader)

model, optimizer = CreateModel(args)

model_D, optimizer_D = CreateDiscriminator(args)

start_iter = 0

if args.restore_from is not None:

start_iter = int(args.restore_from.rsplit('/', 1)[1].rsplit('_')[1])

train_writer = tensorboardX.SummaryWriter(os.path.join(args.snapshot_dir, "logs", model_name))

bce_loss = torch.nn.BCEWithLogitsLoss()

cent_loss=ConditionalEntropyLoss2()

cudnn.enabled = True

cudnn.benchmark = True

model.train()

model.cuda()

model_D.train()

model_D.cuda()

loss = ['loss_seg_src', 'loss_seg_trg', 'loss_D_trg_fake', 'loss_D_src_real', 'loss_D_trg_real']

_t['iter time'].tic()

pbar = tqdm(range(start_iter,args.num_steps_stop))

#for i in range(start_iter, args.num_steps):

for i in pbar:

model.adjust_learning_rate(args, optimizer, i)

model_D.adjust_learning_rate(args, optimizer_D, i)

optimizer.zero_grad()

optimizer_D.zero_grad()

for param in model_D.parameters():

param.requires_grad = False

src_img, src_lbl, _, _ = sourceloader_iter.next()

src_img, src_lbl = Variable(src_img).cuda(), Variable(src_lbl.long()).cuda()

src_seg_score = model(src_img)

loss_seg_src = CrossEntropy2d(src_seg_score, src_lbl)

#loss_seg_src = model.loss

loss_seg_src.backward()

if args.data_label_folder_target is not None:

trg_img, trg_lbl, _, _ = targetloader_iter.next()

trg_img, trg_lbl = Variable(trg_img).cuda(), Variable(trg_lbl.long()).cuda()

trg_seg_score = model(trg_img)

loss_seg_trg = model.loss

else:

trg_img, _, name = targetloader_iter.next()

trg_img = Variable(trg_img).cuda()

trg_seg_score = model(trg_img)

loss_seg_trg=0

#ipdb.set_trace()

outD_trg = model_D(F.softmax(trg_seg_score))

loss_D_trg_fake = bce_loss(outD_trg, Variable(torch.FloatTensor(outD_trg.data.size()).fill_(0)).cuda())

src_seg_score1, trg_seg_score1 = src_seg_score.detach(), trg_seg_score.detach()

if i > args.warm_up:

_, _, h, w = trg_seg_score.size()

outD_trg = nn.functional.upsample(outD_trg, (h, w), mode='bilinear', align_corners=True)

D_out_sigmoid = F.sigmoid(outD_trg).data.cpu().numpy().squeeze(axis=1)

ignore_mask = (D_out_sigmoid > args.mask_T)

loss_seg_trg=cent_loss(trg_seg_score)

ipdb.set_trace()

loss_seg_trg[ignore_mask] = 0

loss_seg_trg=-torch.mean(loss_seg_trg)

#loss_seg_trg = CrossEntropy2d(trg_seg_score, tar_gt)

loss_trg = args.lambda_adv_target * loss_D_trg_fake + args.tar_vat*loss_seg_trg

loss_trg.backward()

if loss_seg_trg ==0:

loss_seg_trg = torch.zeros(1)

for param in model_D.parameters():

param.requires_grad = True

#src_seg_score, trg_seg_score = src_seg_score.detach(), trg_seg_score.detach()

outD_src = model_D(F.softmax(src_seg_score1))

loss_D_src_real = bce_loss(outD_src, Variable(torch.FloatTensor(outD_src.data.size()).fill_(0)).cuda())/ 2

loss_D_src_real.backward()

outD_trg = model_D(F.softmax(trg_seg_score1))

loss_D_trg_real = bce_loss(outD_trg, Variable(torch.FloatTensor(outD_trg.data.size()).fill_(1)).cuda())/ 2

loss_D_trg_real.backward()

d_loss=loss_D_src_real.data + loss_D_trg_real.data

optimizer.step()

optimizer_D.step()

for m in loss:

train_writer.add_scalar(m, eval(m), i+1)

if (i+1) % args.save_pred_every == 0:

print 'taking snapshot ...'

torch.save(model.state_dict(), os.path.join(args.snapshot_dir, '%s_' %(args.source) +str(i+1)+'.pth' ))

torch.save(model_D.state_dict(), os.path.join(args.snapshot_dir, '%s_' %(args.source) +str(i+1)+'_D.pth' ))

if (i+1) % args.print_freq == 0:

_t['iter time'].toc(average=False)

print '[it %d][src seg loss %.4f][trg seg loss %.4f][adv loss %.4f][d loss %.4f][lr %.4f][%.2fs]' % \

(i + 1, loss_seg_src.data,loss_seg_trg.data, loss_D_trg_fake.data,d_loss,optimizer.param_groups[0]['lr']*10000, _t['iter time'].diff)

if i + 1 > args.num_steps_stop:

print 'finish training'

break