def check_original_discriminator(args, pred_target1, pred_target2, n_iter): device = torch.device("cuda" if not args.cpu else "cpu") model_D1 = FCDiscriminator(num_classes=args.num_classes).to(device) model_D2 = FCDiscriminator(num_classes=args.num_classes).to(device) model_D1.eval() model_D2.eval() baseline_dir = 'snapshots/baseline_single_250000_seg0.1_adv10.0002_adv20.001_bs1_11-10-8-52/' bce_loss = torch.nn.BCEWithLogitsLoss() #### restore model_D1, D2 and model # model_D1 parameters D1 = baseline_dir + 'GTA5_150000_D1.pth' saved_state_dict = torch.load(D1) model_D1.load_state_dict(saved_state_dict) # model_D2 parameters D2 = baseline_dir + 'GTA5_150000_D2.pth' saved_state_dict = torch.load(D2) model_D2.load_state_dict(saved_state_dict) D_out1 = model_D1(F.softmax(pred_target1)) D_out2 = model_D2(F.softmax(pred_target2)) orig_d1_loss = bce_loss( D_out1, torch.FloatTensor(D_out1.data.size()).fill_(1).to(device)) orig_d2_loss = bce_loss( D_out2, torch.FloatTensor(D_out2.data.size()).fill_(1).to(device)) original_discriminator_file = args.snapshot_dir + '/original_discriminator.txt' with open(original_discriminator_file, 'a+') as f: f.write('n_iter:{} '.format(n_iter)) f.write('orig d1 loss:{} '.format(orig_d1_loss)) f.write('orig d2 loss:{} '.format(orig_d2_loss)) f.write('\n') del model_D1 del model_D2
def main(): """Create the model and start the training.""" h, w = map(int, args.input_size.split(',')) input_size = (h, w) h, w = map(int, args.input_size_target.split(',')) input_size_target = (h, w) #cudnn.enabled = True gpu = args.gpu # Create network #if args.model == 'DeepLab': # model = DeeplabMulti(num_classes=args.num_classes) # if args.restore_from[:4] == 'http' : # saved_state_dict = model_zoo.load_url(args.restore_from) # else: # saved_state_dict = torch.load(args.restore_from) # new_params = model.state_dict().copy() # for i in saved_state_dict: # # Scale.layer5.conv2d_list.3.weight # i_parts = i.split('.') # # print i_parts # if not args.num_classes == 19 or not i_parts[1] == 'layer5': # new_params['.'.join(i_parts[1:])] = saved_state_dict[i] # print i_parts # model.load_state_dict(new_params) #model.train() #model.cuda(args.gpu) #cudnn.benchmark = True # init D model_D = FCDiscriminator(num_classes=args.num_classes) model_D.train() model_D.cuda(0) #model_D1 = FCDiscriminator(num_classes=args.num_classes) #model_D2 = FCDiscriminator(num_classes=args.num_classes) #model_D1.train() #model_D1.cuda(args.gpu) #model_D2.train() #model_D2.cuda(args.gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) gta_trainloader = data.DataLoader(GTA5DataSet( args.data_dir, args.data_list, args.num_classes, max_iters=args.num_steps * args.iter_size * args.batch_size, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) gta_trainloader_iter = enumerate(gta_trainloader) gta_valloader = data.DataLoader(GTA5DataSet(args.data_dir, args.valdata_list, args.num_classes, max_iters=None, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) gta_valloader_iter = enumerate(gta_valloader) cityscapes_targetloader = data.DataLoader(cityscapesDataSet( args.data_dir_target, args.data_list_target, args.num_classes, max_iters=args.num_steps * args.iter_size * args.batch_size, crop_size=input_size_target, scale=False, mirror=args.random_mirror, mean=IMG_MEAN, set='train'), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) cityscapes_targetloader_iter = enumerate(cityscapes_targetloader) cityscapes_valtargetloader = data.DataLoader(cityscapesDataSet( args.data_dir_target, args.valdata_list_target, args.num_classes, max_iters=None, crop_size=input_size_target, scale=False, mirror=args.random_mirror, mean=IMG_MEAN, set='val'), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) cityscapes_valtargetloader_iter = enumerate(cityscapes_valtargetloader) # implement model.optim_parameters(args) to handle different models' lr setting #optimizer = optim.SGD(model.optim_parameters(args), # lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) #optimizer.zero_grad() optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() #optimizer_D1 = optim.Adam(model_D1.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) #optimizer_D1.zero_grad() #optimizer_D2 = optim.Adam(model_D2.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) #optimizer_D2.zero_grad() bce_loss = torch.nn.BCEWithLogitsLoss() #interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') #interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear') # labels for adversarial training source_label = 0 target_label = 1 for i_iter in range(args.num_steps): loss_D_value = 0 #loss_seg_value1 = 0 #loss_adv_target_value1 = 0 #loss_D_value1 = 0 #loss_seg_value2 = 0 #loss_adv_target_value2 = 0 #loss_D_value2 = 0 optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) _, batch = gta_trainloader_iter.next() images, labels, _, _ = batch size = labels.size() #print(size) #labels = Variable(labels) oneHot_size = (size[0], args.num_classes, size[2], size[3]) input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_() input_label = input_label.scatter_(1, labels.long(), 1.0) #print(input_label.size()) labels1 = Variable(input_label).cuda(0) #D_out1 = model_D(labels) #print(D_out1.data.size()) #loss_out1 = bce_loss(D_out1, Variable(torch.FloatTensor(D_out1.data.size()).fill_(source_label)).cuda(0)) _, batch = cityscapes_targetloader_iter.next() images, labels, _, _ = batch size = labels.size() #labels = Variable(labels) oneHot_size = (size[0], args.num_classes, size[2], size[3]) input_label = torch.FloatTensor(torch.Size(oneHot_size)).zero_() input_label = input_label.scatter_(1, labels.long(), 1.0) labels2 = Variable(input_label).cuda(0) #print(labels1.data.size()) #print(labels2.data.size()) labels = torch.cat((labels1, labels2), 0) #print(labels.data.size()) #D_out2 = model_D(labels) D_out = model_D(labels) #print(D_out.data.size()) target_size = D_out.data.size() target_labels1 = torch.FloatTensor( torch.Size((target_size[0] / 2, target_size[1], target_size[2], target_size[3]))).fill_(source_label) target_labels2 = torch.FloatTensor( torch.Size((target_size[0] / 2, target_size[1], target_size[2], target_size[3]))).fill_(target_label) target_labels = torch.cat((target_labels1, target_labels2), 0) target_labels = Variable(target_labels).cuda(0) #print(target_labels.data.size()) loss_out = bce_loss(D_out, target_labels) loss = loss_out / args.iter_size loss.backward() loss_D_value += loss_out.data.cpu().numpy()[0] / args.iter_size #print(loss_D_value) optimizer_D.step() #print('exp = {}'.format(args.snapshot_dir)) if i_iter % 100 == 0: print('iter = {0:8d}/{1:8d}, loss_D = {2:.3f}'.format( i_iter, args.num_steps, loss_D_value)) if i_iter >= args.num_steps_stop - 1: print('save model ...') torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'Classify_' + str(args.num_steps) + '.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'Classify_' + str(i_iter) + '.pth')) model_D.eval() loss_valD_value = 0 correct = 0 wrong = 0 for i, (images, labels, _, _) in enumerate(gta_valloader): #if i > 500: # break size = labels.size() #labels = Variable(labels) oneHot_size = (size[0], args.num_classes, size[2], size[3]) input_label = torch.FloatTensor( torch.Size(oneHot_size)).zero_() input_label = input_label.scatter_(1, labels.long(), 1.0) labels = Variable(input_label).cuda(0) D_out1 = model_D(labels) loss_out1 = bce_loss( D_out1, Variable( torch.FloatTensor( D_out1.data.size()).fill_(source_label)).cuda(0)) loss_valD_value += loss_out1.data.cpu().numpy()[0] correct = correct + (D_out1.data.cpu() < 0).sum() / 100 wrong = wrong + (D_out1.data.cpu() >= 0).sum() / 100 #accuracy = 1.0 * correct / (wrong + correct) #print('accuracy:%f' % accuracy) #print(correct) #print(wrong) for i, (images, labels, _, _) in enumerate(cityscapes_valtargetloader): #if i > 500: # break size = labels.size() #labels = Variable(labels) oneHot_size = (size[0], args.num_classes, size[2], size[3]) input_label = torch.FloatTensor( torch.Size(oneHot_size)).zero_() input_label = input_label.scatter_(1, labels.long(), 1.0) labels = Variable(input_label).cuda(0) D_out2 = model_D(labels) loss_out2 = bce_loss( D_out2, Variable( torch.FloatTensor( D_out2.data.size()).fill_(target_label)).cuda(0)) loss_valD_value += loss_out2.data.cpu().numpy()[0] wrong = wrong + (D_out2.data.cpu() < 0).sum() correct = correct + (D_out2.data.cpu() >= 0).sum() accuracy = 1.0 * correct / (wrong + correct) print('accuracy:%f' % accuracy) print('iter = {0:8d}/{1:8d}, loss_valD = {2:.3f}'.format( i_iter, args.num_steps, loss_valD_value)) model_D.train()
def main(): """Create the model and start the evaluation process.""" args = get_arguments() if not os.path.exists(args.save): os.makedirs(args.save) gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) if args.restore_from[:4] == 'http': saved_state_dict = model_zoo.load_url(args.restore_from) else: saved_state_dict = torch.load(args.restore_from) model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) #===========load discriminator model====begin=== model_d2 = FCDiscriminator(num_classes=args.num_classes) d2_state_dict = torch.load(args.dis_restore_from) model_d2.load_state_dict(d2_state_dict) model_d2.eval() model_d2.cuda(gpu0) #===========load discriminator model====end=== testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set), batch_size=1, shuffle=False, pin_memory=True) interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True) # interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear') out_values = [] fine_out_values = [] retrain_list = [] file = open(CITYS_RETRAIN_TXT, 'w') for index, batch in enumerate(testloader): if index % 20 == 0: print('%d processd of %d' % (index, len(testloader))) image, _, name = batch output1, output2 = model(Variable(image, volatile=True).cuda(gpu0)) ini_output = interp(output2) d2_out1 = model_d2(F.softmax(ini_output, dim=1)) #.cpu().data[0].numpy() out_valu = d2_out1.mean() out_valu_img = np.array([[name[0]], out_valu.cpu().data.numpy()]) out_values.extend(out_valu_img) if out_valu.cpu().data.numpy() > 0.64: fine_out_valu_img = np.array([[name[0]], out_valu.cpu().data.numpy()]) fine_out_values.extend(fine_out_valu_img) file.write(name[0] + '\n') output = interp(output2).cpu().data[0].numpy() output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) output_col = colorize_mask(output) name = name[0].split('/')[-1] # output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0])) output_col.save('%s/%s.png' % (args.save, name.split('.')[0])) # print('its confidence value is %f' % out_valu) # plt.imshow(output_col) # plt.title(str(out_valu)) # plt.show() # output = Image.fromarray(output) # output.save('%s/%s' % (args.save, name)) out_values = np.array(out_values) np.save(CITYS_VALUES_SV_PATH, out_values) np.save(CITYS_FINE_VALUES_SV_PATH, fine_out_values) file.close()