def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) model = Res_Deeplab(num_classes=args.num_classes) if args.pretrained_model != None: args.restore_from = pretrianed_models_dict[args.pretrained_model] 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) testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] colorize = VOCColorize() for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd'%(index)) image, label, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output).cpu().data[0].numpy() output = output[:,:size[0],:size[1]] gt = np.asarray(label[0].numpy()[:size[0],:size[1]], dtype=np.int) output = output.transpose(1,2,0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) filename = os.path.join(args.save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray(colorize(output).transpose(1, 2, 0), 'RGB') color_file.save(filename) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) filename = os.path.join(args.save_dir, 'result.txt') get_iou(data_list, args.num_classes, filename)
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) model = Res_Deeplab(num_classes=args.num_classes) if args.pretrained_model != None: args.restore_from[i] = pretrianed_models_dict[args.pretrained_model] # if there is pretrained model, restore_from will be changed if args.restore_from[i][:4] == 'http' : saved_state_dict = model_zoo.load_url(args.restore_from[i]) else: saved_state_dict = torch.load(args.restore_from[i]) ##VOC_25000 model.load_state_dict(saved_state_dict) model.eval() #evaluation mode model.cuda(gpu0) testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] colorize = VOCColorize() # colorize! with torch.no_grad():### added for 0.4 for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd'%(index)) image, label, size, name = batch # size >> tensor([[366, 500, 3]]) size = size[0].numpy() ## [366 500 3] output = model(image.cuda(args.gpu)) output = interp(output).cpu().data[0].numpy()# 21,505,505 output = output[:,:size[0],:size[1]] # 21,366,500 gt = np.asarray(label[0].numpy()[:size[0],:size[1]], dtype=np.int) # size of each image is diff output = output.transpose(1,2,0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) filename = os.path.join(args.save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray(colorize(output).transpose(1, 2, 0), 'RGB') # colorize the output color_file.save(filename) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) filename = os.path.join(args.save_dir, 'result'+args.restore_from[i][-10:-4]+'.txt') get_iou(data_list, args.num_classes, filename)
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if not os.path.exists(args.save): os.makedirs(args.save) if args.model == 'Deeplab': model = Res_Deeplab(num_classes=args.num_classes) elif args.model == 'DeeplabVGG': model = DeeplabVGG(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from) ### for running different versions of pytorch model_dict = model.state_dict() saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict} model_dict.update(saved_state_dict) ### model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) 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) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(1024, 2048), mode='bilinear') for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd' % index) image, _, name = batch if args.model == 'Deeplab': output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output).cpu().data[0].numpy() elif args.model == 'DeeplabVGG': output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output).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) output = Image.fromarray(output) name = name[0].split('/')[-1] output.save('%s/%s' % (args.save, name)) output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
def eval(pth, cityscapes_eval_dir, i_iter): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if args.model == 'ResNet': model = Res_Deeplab(num_classes=args.num_classes) saved_state_dict = torch.load(pth) elif args.model == 'VGG': model = DeeplabVGG(num_classes=args.num_classes) saved_state_dict = torch.load(pth) model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) cityscapesloader = data.DataLoader(cityscapesDataSet( args.cityscapes_data_dir, args.cityscapes_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) for index, batch in enumerate(cityscapesloader): with torch.no_grad(): if index % 100 == 0: print('%d processd' % index) image, _, name = batch output = model(Variable(image).cuda(gpu0)) output = interp(output).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) output = Image.fromarray(output) name = name[0].split('/')[-1] output.save('%s/%s' % (cityscapes_eval_dir, name)) output_col.save('%s/%s_color.png' % (cityscapes_eval_dir, name.split('.')[0])) if i_iter == 0: break
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) model = Res_Deeplab(num_classes=args.num_classes) # if args.pretrained_model != None: # args.restore_from = pretrianed_models_dict[args.pretrained_model] # # 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 = Res_Deeplab(num_classes=args.num_classes) #model.load_state_dict(torch.load('/data/wyc/AdvSemiSeg/snapshots/VOC_15000.pth')) state_dict=torch.load('/data1/wyc/AdvSemiSeg/snapshots/VOC_t_concat_pred_img_15000.pth') from model.discriminator_pred_concat_img import FCDiscriminator model_D = FCDiscriminator(num_classes=args.num_classes) state_dict_d = torch.load('/data1/wyc/AdvSemiSeg/snapshots/VOC_t_concat_pred_img_15000_D.pth') # original saved file with DataParallel # create new OrderedDict that does not contain `module.` from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v # load params new_params = model.state_dict().copy() for name, param in new_params.items(): print (name) if name in new_state_dict and param.size() == new_state_dict[name].size(): new_params[name].copy_(new_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.eval() model.cuda() new_state_dict_d = OrderedDict() for k, v in state_dict_d.items(): name = k[7:] # remove `module.` new_state_dict_d[name] = v new_params_d = model_D.state_dict().copy() for name, param in new_params_d.items(): print (name) if name in new_state_dict_d and param.size() == new_state_dict_d[name].size(): new_params_d[name].copy_(new_state_dict_d[name]) print('copy {}'.format(name)) model_D.load_state_dict(new_params_d) model_D.eval() model_D.cuda() testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] colorize = VOCColorize() for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd'%(index)) image, label, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda()) image_d=Variable(image, volatile=True).cuda() output=interp(output) output_dout = output.clone() output_pred = F.softmax(output, dim=1).cpu().data[0].numpy() label2=label[0].numpy() output = output.cpu().data[0].numpy() output = output[:,:size[0],:size[1]] gt = np.asarray(label[0].numpy()[:size[0],:size[1]], dtype=np.int) output = output.transpose(1,2,0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) #"""pred result""" filename = os.path.join(args.save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray(colorize(output).transpose(1, 2, 0), 'RGB') color_file.save(filename) #"""the area of the pred which is wrong""" output_mistake=np.zeros(output.shape) semi_ignore_mask_correct = (output == gt) #semi_ignore_mask_255=(gt==255) output_mistake[semi_ignore_mask_correct] = 255 #output_mistake[semi_ignore_mask_255] = 255 output_mistake = np.expand_dims(output_mistake, axis=2) filename2 = os.path.join('/data1/wyc/AdvSemiSeg/pred_mis/', '{}.png'.format(name[0])) cv2.imwrite(filename2, output_mistake) #"""dis confidence map decide line of pred map""" D_out = interp(model_D(torch.cat([F.softmax(output_dout, dim=1),F.sigmoid(image_d)],1)))#67 D_out_sigmoid = (F.sigmoid(D_out).data[0].cpu().numpy()) D_out_sigmoid = D_out_sigmoid[:, :size[0], :size[1]] semi_ignore_mask_dout0 = (D_out_sigmoid < 0.0001) semi_ignore_mask_dout255 = (D_out_sigmoid >= 0.0001) D_out_sigmoid[semi_ignore_mask_dout0] = 0 D_out_sigmoid[semi_ignore_mask_dout255] = 255 filename2 = os.path.join('/data1/wyc/AdvSemiSeg/confidence_line/', '{}.png'.format(name[0]))#0 black 255 white cv2.imwrite(filename2,D_out_sigmoid.transpose(1, 2, 0)) #""""pred max decide line of pred map""" # id2 = np.argmax(output_pred, axis=0) # map=np.zeros([1,id2.shape[0],id2.shape[1]]) # for i in range(id2.shape[0]): # for j in range(id2.shape[1]): # map[0][i][j]=output_pred[id2[i][j]][i][j] # semi_ignore_mask2 = (map < 0.999999) # semi_ignore_mask3 = (map >= 0.999999) # map[semi_ignore_mask2] = 0 # map[semi_ignore_mask3] = 255 # map = map[:, :size[0], :size[1]] # filename2 = os.path.join('/data1/wyc/AdvSemiSeg/pred_line/', '{}.png'.format(name[0]))#0 black 255 white # cv2.imwrite(filename2,map.transpose(1, 2, 0)) data_list.append([gt.flatten(), output.flatten()]) filename = os.path.join(args.save_dir, 'result.txt') get_iou(data_list, args.num_classes, filename)
def main(): # 将参数的input_size 映射到整数,并赋值,从字符串转换到整数二元组 h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = False gpu = args.gpu # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) # 确保模型中参数的格式与要加载的参数相同 # 返回一个字典,保存着module的所有状态(state);parameters和persistent buffers都会包含在字典中,字典的key就是parameter和buffer的 names。 new_params = model.state_dict().copy() for name, param in new_params.items(): # print (name) if name in saved_state_dict and param.size() == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) # print('copy {}'.format(name)) model.load_state_dict(new_params) # 设置为训练模式 model.train() cudnn.benchmark = True model.cuda(gpu) # init D model_D = FCDiscriminator(num_classes=args.num_classes) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: # sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = list(range(train_dataset_size)) # ? np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) # 写入文件 train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # ??? # labels for adversarial training pred_label = 0 gt_label = 1 for i_iter in range(args.num_steps): print("Iter:", i_iter) loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_semi_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): # train G # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # do semi first if args.lambda_semi > 0 and i_iter >= args.semi_start: try: _, batch = next(trainloader_remain_iter) except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = next(trainloader_remain_iter) # only access to img images, _, _, _ = batch images = Variable(images).cuda(gpu) # images = Variable(images).cpu() pred = interp(model(images)) D_out = interp(model_D(F.softmax(pred))) D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1) # produce ignore mask semi_ignore_mask = (D_out_sigmoid < args.mask_T) semi_gt = pred.data.cpu().numpy().argmax(axis=1) semi_gt[semi_ignore_mask] = 255 semi_ratio = 1.0 - float(semi_ignore_mask.sum()) / semi_ignore_mask.size print('semi ratio: {:.4f}'.format(semi_ratio)) if semi_ratio == 0.0: loss_semi_value += 0 else: semi_gt = torch.FloatTensor(semi_gt) loss_semi = args.lambda_semi * loss_calc(pred, semi_gt, args.gpu) loss_semi = loss_semi / args.iter_size loss_semi.backward() loss_semi_value += loss_semi.data.cpu().numpy()[0] / args.lambda_semi else: loss_semi = None # train with source try: _, batch = next(trainloader_iter) except: trainloader_iter = enumerate(trainloader) _, batch = next(trainloader_iter) images, labels, _, _ = batch images = Variable(images).cuda(gpu) # images = Variable(images).cpu() ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) loss_seg = loss_calc(pred, labels, args.gpu) D_out = interp(model_D(F.softmax(pred))) loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask)) loss = loss_seg + args.lambda_adv_pred * loss_adv_pred # proper normalization loss = loss / args.iter_size loss.backward() loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0] / args.iter_size # train D # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred pred = pred.detach() D_out = interp(model_D(F.softmax(pred))) loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] # train with gt # get gt labels try: _, batch = next(trainloader_gt_iter) except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = next(trainloader_gt_iter) _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) # D_gt_v = Variable(one_hot(labels_gt)).cpu() ignore_mask_gt = (labels_gt.numpy() == 255) D_out = interp(model_D(D_gt_v)) loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] optimizer.step() optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) print( 'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}'.format( i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value)) if i_iter >= args.num_steps - 1: print('save model ...') torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '.pth')) torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth')) torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): # prepare h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True gpu = args.gpu if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() for name, param in new_params.items(): if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.train() model.cuda(args.gpu) cudnn.benchmark = True # init D model_D = detector.FlawDetector(in_channels=24) # model_D = FCDiscriminator(num_classes=args.num_classes) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(args.gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: #sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = [_ for _ in range(0, train_dataset_size)] np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) train_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler( train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() # loss/ bilinear upsampling minimum_loss = detector.MinimumCriterion() detector_loss = detector.FlawDetectorCriterion() # bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 for i_iter in range(args.num_steps): if i_iter > 0 and i_iter % 1000 == 0: val(model, args.gpu) model.train() loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_semi_value = 0 loss_semi_adv_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): # train G # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # do semi first if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv: try: _, batch = trainloader_remain_iter.__next__() except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = trainloader_remain_iter.__next__() # only access to img images, _, _, _ = batch images = Variable(images).cuda(args.gpu) pred = interp(model(images)) pred_remain = pred.detach() D_out = model_D(images, pred) ignore_mask_remain = np.zeros(D_out.shape).astype(np.bool) loss_semi_adv = args.lambda_semi_adv * minimum_loss( D_out) # ke: SSL loss for unlabeled data loss_semi_adv = loss_semi_adv #loss_semi_adv.backward() loss_semi_adv_value += loss_semi_adv.data.cpu().numpy( ) / args.iter_size loss_semi_adv.backward() loss_semi_value = 0 else: loss_semi = None loss_semi_adv = None # train with source (labeled data) try: _, batch = trainloader_iter.__next__() except: trainloader_iter = enumerate(trainloader) _, batch = trainloader_iter.__next__() images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) loss_seg = loss_calc(pred, labels, args.gpu) D_out = interp(model_D(images, pred)) loss_adv_pred = args.lambda_adv_pred * minimum_loss( D_out) # ke: SSL loss for labeled data loss = loss_seg + loss_adv_pred # proper normalization loss = loss / args.iter_size loss.backward() loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size loss_adv_pred_value += loss_adv_pred.data.cpu().numpy( ) / args.iter_size # train D # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred from labeled data pred = pred.detach() D_out = interp(model_D(images, pred)) detect_gt = detector.generate_flaw_detector_gt( pred, labels.view(labels.shape[0], 1, labels.shape[1], labels.shape[2]).cuda(args.gpu), NUM_CLASSES, IGNORE_LABEL) loss_D = detector_loss(D_out, detect_gt) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy() # # train with gt # # get gt labels # try: # _, batch = trainloader_gt_iter.__next__() # except: # trainloader_gt_iter = enumerate(trainloader_gt) # _, batch = trainloader_gt_iter.__next__() # _, labels_gt, _, _ = batch # D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) # ignore_mask_gt = (labels_gt.numpy() == 255) # D_out = interp(model_D(D_gt_v)) # loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) # loss_D = loss_D/args.iter_size/2 # loss_D.backward() # loss_D_value += loss_D.data.cpu().numpy() optimizer.step() optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) print( 'iter = {0:8d}/{1:8d}, loss_seg = {2:.6f}, loss_adv_l = {3:.6f}, loss_D = {4:.6f}, loss_semi = {5:.6f}, loss_adv_u = {6:.6f}' .format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value, loss_semi_adv_value)) if i_iter >= args.num_steps - 1: torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth')) end = timeit.default_timer()
def main(): h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True gpu = args.gpu # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters (weights) if args.restore_from[:4] == 'http': saved_state_dict = model_zoo.load_url( args.restore_from ) ## http://vllab1.ucmerced.edu/~whung/adv-semi-seg/resnet101COCO-41f33a49.pth else: saved_state_dict = torch.load(args.restore_from) #checkpoint = torch.load(args.restore_from)_ # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() # state_dict() is current model for name, param in new_params.items(): #print (name) # 'conv1.weight, name:param(value), dict if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) #print('copy {}'.format(name)) model.load_state_dict(new_params) #model.load_state_dict(checkpoint['state_dict']) #optimizer.load_state_dict(args.checkpoint['optim_dict']) model.train( ) # https://pytorch.org/docs/stable/nn.html, Sets the module in training mode. model.cuda(args.gpu) ## cudnn.benchmark = True # This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware # init D model_D = FCDiscriminator(num_classes=args.num_classes) #args.restore_from_D = 'snapshots/linear2/VOC_25000_D.pth' if args.restore_from_D is not None: # None model_D.load_state_dict(torch.load(args.restore_from_D)) # checkpoint_D = torch.load(args.restore_from_D) # model_D.load_state_dict(checkpoint_D['state_dict']) # optimizer_D.load_state_dict(checkpoint_D['optim_dict']) model_D.train() model_D.cuda(args.gpu) if USECALI: model_cali = ModelWithTemperature(model, model_D) model_cali.cuda(args.gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) random.seed(args.random_seed) np.random.seed(args.random_seed) torch.manual_seed(args.random_seed) torch.cuda.manual_seed(args.random_seed) train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_remain = VOCDataSet(args.data_dir, args.data_list_remain, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_dataset_size_remain = len(train_dataset_remain) print train_dataset_size print train_dataset_size_remain train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: #if not partial, load all trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: #sample partial data #args.partial_data = 0.125 partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: #args.partial_id is none train_ids = range(train_dataset_size) train_ids_remain = range(train_dataset_size_remain) np.random.shuffle(train_ids) #shuffle! np.random.shuffle(train_ids_remain) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) #randomly suffled ids #sampler train_sampler = data.sampler.SubsetRandomSampler( train_ids[:]) # 0~1/8, train_remain_sampler = data.sampler.SubsetRandomSampler( train_ids_remain[:]) train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:]) # train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) # 0~1/8 # train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:]) # used as unlabeled, 7/8 # train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) #train loader trainloader = data.DataLoader( train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) # multi-process data loading trainloader_remain = data.DataLoader(train_dataset_remain, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) # trainloader_remain = data.DataLoader(train_dataset, # batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, # pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network # model.optim_paramters(args) = list(dict1, dict2), dict1 >> 'lr' and 'params' # print(type(model.optim_parameters(args)[0]['params'])) # generator #print(model.state_dict()['coeff'][0]) #confirmed optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) #optimizer.add_param_group({"params":model.coeff}) # assign new coefficient to the optimizer #print(len(optimizer.param_groups)) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() #initialize if USECALI: optimizer_cali = optim.LBFGS([model_cali.temperature], lr=0.01, max_iter=50) optimizer_cali.zero_grad() nll_criterion = BCEWithLogitsLoss().cuda() # BCE!! ece_criterion = ECELoss().cuda() # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample( size=(input_size[1], input_size[0]), mode='bilinear' ) # okay it automatically change to functional.interpolate # 321, 321 if version.parse(torch.__version__) >= version.parse('0.4.0'): #0.4.1 interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 semi_ratio_sum = 0 semi_sum = 0 loss_seg_sum = 0 loss_adv_sum = 0 loss_vat_sum = 0 l_seg_sum = 0 l_vat_sum = 0 l_adv_sum = 0 logits_list = [] labels_list = [] #https: // towardsdatascience.com / understanding - pytorch -with-an - example - a - step - by - step - tutorial - 81fc5f8c4e8e for i_iter in range(args.num_steps): loss_seg_value = 0 # L_seg loss_adv_pred_value = 0 # 0.01 L_adv loss_D_value = 0 # L_D loss_semi_value = 0 # 0.1 L_semi loss_semi_adv_value = 0 # 0.001 L_adv loss_vat_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) #changing lr by iteration optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): ###################### train G!!!########################### ############################################################ # don't accumulate grads in D for param in model_D.parameters( ): # <class 'torch.nn.parameter.Parameter'>, convolution weights param.requires_grad = False # do not update gradient of D (freeze) while G ######### do unlabeled first!! 0.001 L_adv + 0.1 L_semi ############### # lambda_semi, lambda_adv for unlabeled if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv: try: _, batch = trainloader_remain_iter.next( ) #remain = unlabeled print(trainloader_remain_iter.next()) except: trainloader_remain_iter = enumerate( trainloader_remain) # impose counters _, batch = trainloader_remain_iter.next() # only access to img images, _, _, _ = batch # <class 'torch.Tensor'> images = Variable(images).cuda( args.gpu) # <class 'torch.Tensor'> pred = interp( model(images)) # S(X), pred <class 'torch.Tensor'> pred_remain = pred.detach( ) #use detach() when attempting to remove a tensor from a computation graph, will be used for D # https://discuss.pytorch.org/t/clone-and-detach-in-v0-4-0/16861 # The difference is that detach refers to only a given variable on which it's called. # torch.no_grad affects all operations taking place within the with statement. >> for context, # requires_grad is for tensor # pred >> (8,21,321,321), L_adv D_out = interp( model_D(F.softmax(pred)) ) # D(S(X)), confidence, 8,1,321,321, not detached, there was not dim D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze( axis=1) # (8,321,321) 0~1 # 0.001 L_adv!!!! ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype( np.bool) # no ignore_mask for unlabeled adv loss_semi_adv = args.lambda_semi_adv * bce_loss( D_out, make_D_label(gt_label, ignore_mask_remain)) #gt_label =1, # -log(D(S(X))) loss_semi_adv = loss_semi_adv / args.iter_size #normalization loss_semi_adv_value += loss_semi_adv.data.cpu().numpy( ) / args.lambda_semi_adv ##--- visualization, pred(8,21,321,321), D_out_sigmoid(8,321,321) """ if i_iter % 1000 == 0: vpred = pred.transpose(1, 2).transpose(2, 3).contiguous() # (8,321,321,21) vpred = vpred.view(-1, 21) # (8*321*321, 21) vlogsx = F.log_softmax(vpred) # torch.Tensor vsemi_gt = pred.data.cpu().numpy().argmax(axis=1) vsemi_gt = Variable(torch.FloatTensor(vsemi_gt).long()).cuda(gpu) vlogsx = vlogsx.gather(1, vsemi_gt.view(-1, 1)) sx = F.softmax(vpred).gather(1, vsemi_gt.view(-1, 1)) vD_out_sigmoid = Variable(torch.FloatTensor(D_out_sigmoid)).cuda(gpu).view(-1, 1) vlogsx = (vlogsx*(2.5*vD_out_sigmoid+0.5)) vlogsx = -vlogsx.squeeze(dim=1) sx = sx.squeeze(dim=1) vD_out_sigmoid = vD_out_sigmoid.squeeze(dim=1) dsx = vD_out_sigmoid.data.cpu().detach().numpy() vlogsx = vlogsx.data.cpu().detach().numpy() sx = sx.data.cpu().detach().numpy() plt.clf() plt.figure(figsize=(15, 5)) plt.subplot(131) plt.ylim(0, 0.004) plt.scatter(dsx, vlogsx, s = 0.1) # variable requires grad cannot call numpy >> detach plt.xlabel('D(S(X))') plt.ylabel('Loss_Semi per Pixel') plt.subplot(132) plt.scatter(dsx, vlogsx, s = 0.1) # variable requires grad cannot call numpy >> detach plt.xlabel('D(S(X))') plt.ylabel('Loss_Semi per Pixel') plt.subplot(133) plt.scatter(dsx, sx, s=0.1) plt.xlabel('D(S(X))') plt.ylabel('S(x)') plt.savefig('/home/eungyo/AdvSemiSeg/plot/' + str(i_iter) + '.png') """ if args.lambda_semi <= 0 or i_iter < args.semi_start: loss_semi_adv.backward() loss_semi_value = 0 else: semi_gt = pred.data.cpu().numpy().argmax( axis=1 ) # pred=S(X) ((8,21,321,321)), semi_gt is not one-hot, 8,321,321 #(8, 321, 321) if not USECALI: semi_ignore_mask = ( D_out_sigmoid < args.mask_T ) # both (8,321,321) 0~1threshold!, numpy semi_gt[ semi_ignore_mask] = 255 # Yhat, ignore pixel becomes 255 semi_ratio = 1.0 - float(semi_ignore_mask.sum( )) / semi_ignore_mask.size # ignored pixels / H*W print('semi ratio: {:.4f}'.format(semi_ratio)) if semi_ratio == 0.0: loss_semi_value += 0 else: semi_gt = torch.FloatTensor(semi_gt) confidence = torch.FloatTensor( D_out_sigmoid) ## added, only pred is on cuda loss_semi = args.lambda_semi * weighted_loss_calc( pred, semi_gt, args.gpu, confidence) else: semi_ratio = 1 semi_gt = (torch.FloatTensor(semi_gt)) # (8,321,321) confidence = torch.FloatTensor( F.sigmoid( model_cali.temperature_scale(D_out.view( -1))).data.cpu().numpy()) # (8*321*321,) loss_semi = args.lambda_semi * calibrated_loss_calc( pred, semi_gt, args.gpu, confidence, accuracies, n_bin ) # L_semi = Yhat * log(S(X)) # loss_calc(pred, semi_gt, args.gpu) # pred(8,21,321,321) if semi_ratio != 0: loss_semi = loss_semi / args.iter_size loss_semi_value += loss_semi.data.cpu().numpy( ) / args.lambda_semi if args.method == 'vatent' or args.method == 'vat': #v_loss = vat_loss(model, images, pred, eps=args.epsilon[i]) # R_vadv weighted_v_loss = weighted_vat_loss( model, images, pred, confidence, eps=args.epsilon) if args.method == 'vatent': #v_loss += entropy_loss(pred) # R_cent (conditional entropy loss) weighted_v_loss += weighted_entropy_loss( pred, confidence) v_loss = weighted_v_loss / args.iter_size loss_vat_value += v_loss.data.cpu().numpy() loss_semi_adv += args.alpha * v_loss loss_vat_sum += loss_vat_value if i_iter % 100 == 0 and sub_i == 4: l_vat_sum = loss_vat_sum / 100 if i_iter == 0: l_vat_sum = l_vat_sum * 100 loss_vat_sum = 0 loss_semi += loss_semi_adv loss_semi.backward( ) # 0.001 L_adv + 0.1 L_semi, backward == back propagation else: loss_semi = None loss_semi_adv = None ###########train with source (labeled data)############### L_ce + 0.01 * L_adv try: _, batch = trainloader_iter.next() except: trainloader_iter = enumerate(trainloader) # safe coding _, batch = trainloader_iter.next() #counter, batch images, labels, _, _ = batch # also get labels images(8,321,321) images = Variable(images).cuda(args.gpu) ignore_mask = ( labels.numpy() == 255 ) # ignored pixels == 255 >> 1, yes ignored mask for labeled data pred = interp(model(images)) # S(X), 8,21,321,321 loss_seg = loss_calc(pred, labels, args.gpu) # -Y*logS(X)= L_ce, not detached if USED: softsx = F.softmax(pred, dim=1) D_out = interp(model_D(softsx)) # D(S(X)), L_adv loss_adv_pred = bce_loss( D_out, make_D_label( gt_label, ignore_mask)) # both 8,1,321,321, gt_label = 1 # L_adv = -log(D(S(X)), make_D_label is all 1 except ignored_region loss = loss_seg + args.lambda_adv_pred * loss_adv_pred if USECALI: if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv: with torch.no_grad(): _, prediction = torch.max(softsx, 1) labels_mask = ( (labels > 0) * (labels != 255)) | (prediction.data.cpu() > 0) labels = labels[labels_mask] prediction = prediction[labels_mask] fake_mask = (labels.data.cpu().numpy() != prediction.data.cpu().numpy()) real_label = make_conf_label( 1, fake_mask ) # (10*321*321, ) 0 or 1 (fake or real) logits = D_out.squeeze(dim=1) logits = logits[labels_mask] logits_list.append(logits) # initialize labels_list.append(real_label) if (i_iter * args.iter_size * args.batch_size + sub_i + 1) % train_dataset_size == 0: logits = torch.cat(logits_list).cuda( ) # overall 5000 images in val, #logits >> 5000,100, (1464*321*321,) labels = torch.cat(labels_list).cuda() before_temperature_nll = nll_criterion( logits, labels).item() ####modify before_temperature_ece, _, _ = ece_criterion( logits, labels) # (1464*321*321,) before_temperature_ece = before_temperature_ece.item( ) print('Before temperature - NLL: %.3f, ECE: %.3f' % (before_temperature_nll, before_temperature_ece)) def eval(): loss_cali = nll_criterion( model_cali.temperature_scale(logits), labels) loss_cali.backward() return loss_cali optimizer_cali.step( eval) # just one backward >> not 50 iterations after_temperature_nll = nll_criterion( model_cali.temperature_scale(logits), labels).item() after_temperature_ece, accuracies, n_bin = ece_criterion( model_cali.temperature_scale(logits), labels) after_temperature_ece = after_temperature_ece.item( ) print('Optimal temperature: %.3f' % model_cali.temperature.item()) print( 'After temperature - NLL: %.3f, ECE: %.3f' % (after_temperature_nll, after_temperature_ece)) logits_list = [] labels_list = [] else: loss = loss_seg # proper normalization loss = loss / args.iter_size loss.backward() loss_seg_sum += loss_seg / args.iter_size if USED: loss_adv_sum += loss_adv_pred if i_iter % 100 == 0 and sub_i == 4: l_seg_sum = loss_seg_sum / 100 if USED: l_adv_sum = loss_adv_sum / 100 if i_iter == 0: l_seg_sum = l_seg_sum * 100 l_adv_sum = l_adv_sum * 100 loss_seg_sum = 0 loss_adv_sum = 0 loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size if USED: loss_adv_pred_value += loss_adv_pred.data.cpu().numpy( ) / args.iter_size ##################### train D!!!########################### ########################################################### # bring back requires_grad if USED: for param in model_D.parameters(): param.requires_grad = True # before False. ############# train with pred S(X)############# labeled + unlabeled pred = pred.detach( ) #orginally only use labeled data, freeze S(X) when train D, # We do train D with the unlabeled data. But the difference is quite small if args.D_remain: #default true pred = torch.cat( (pred, pred_remain), 0 ) # pred_remain(unlabeled S(x)) is detached 16,21,321,321 ignore_mask = np.concatenate( (ignore_mask, ignore_mask_remain), axis=0) # 16,321,321 D_out = interp( model_D(F.softmax(pred, dim=1)) ) # D(S(X)) 16,1,321,321 # softmax(pred,dim=1) for 0.4, not nessesary loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask)) # pred_label = 0 # -log(1-D(S(X))) loss_D = loss_D / args.iter_size / 2 # iter_size = 1, /2 because there is G and D loss_D.backward() loss_D_value += loss_D.data.cpu().numpy() ################## train with gt################### only labeled #VOCGT and VOCdataset can be reduced to one dataset in this repo. # get gt labels Y #print "before train gt" try: print(trainloader_gt_iter.next()) # len 732 _, batch = trainloader_gt_iter.next() except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = trainloader_gt_iter.next() #print "train with gt?" _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) #one_hot ignore_mask_gt = (labels_gt.numpy() == 255 ) # same as ignore_mask (8,321,321) #print "finish" D_out = interp(model_D(D_gt_v)) # D(Y) loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) # log(D(Y)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy() optimizer.step() if USED: optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) #snapshot print( 'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.6f}, loss_semi = {5:.6f}, loss_semi_adv = {6:.3f}, loss_vat = {7: .5f}' .format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value, loss_semi_adv_value, loss_vat_value)) # L_ce L_adv for labeled L_D L_semi L_adv for unlabeled #loss_adv should be inversely proportional to the loss_D if they are seeing the same data. # loss_adv_p is essentially the inverse loss of loss_D. We expect them to achieve a good balance during the adversarial training # loss_D is around 0.2-0.5 >> good if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_D.pth')) #torch.save(state, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth.tar')) #torch.save(state_D, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth.tar')) break if i_iter % 100 == 0 and sub_i == 4: #loss_seg_value wdata = "iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.6f}, loss_semi = {5:.8f}, loss_semi_adv = {6:.3f}, l_vat_sum = {7: .5f}, loss_label = {8: .4}\n".format( i_iter, args.num_steps, l_seg_sum, l_adv_sum, loss_D_value, loss_semi_value, loss_semi_adv_value, l_vat_sum, l_seg_sum + 0.01 * l_adv_sum) #wdata2 = "{0:8d} {1:s} {2:s} {3:s} {4:s} {5:s} {6:s} {7:s} {8:s}\n".format(i_iter,str(model.coeff[0])[8:14],str(model.coeff[1])[8:14],str(model.coeff[2])[8:14],str(model.coeff[3])[8:14],str(model.coeff[4])[8:14],str(model.coeff[5])[8:14],str(model.coeff[6])[8:14],str(model.coeff[7])[8:14]) if i_iter == 0: f2 = open("/home/eungyo/AdvSemiSeg/snapshots/log.txt", 'w') f2.write(wdata) f2.close() #f3 = open("/home/eungyo/AdvSemiSeg/snapshots/coeff.txt", 'w') #f3.write(wdata2) #f3.close() else: f1 = open("/home/eungyo/AdvSemiSeg/snapshots/log.txt", 'a') f1.write(wdata) f1.close() #f4 = open("/home/eungyo/AdvSemiSeg/snapshots/coeff.txt", 'a') #f4.write(wdata2) #f4.close() if i_iter % args.save_pred_every == 0 and i_iter != 0: # 5000 print('taking snapshot ...') #state = {'epoch':i_iter, 'state_dict':model.state_dict(),'optim_dict':optimizer.state_dict()} #state_D = {'epoch':i_iter, 'state_dict': model_D.state_dict(), 'optim_dict': optimizer_D.state_dict()} #torch.save(state, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth.tar')) #torch.save(state_D, osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth.tar')) torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True gpu = args.gpu # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() for name, param in new_params.items(): print(name) if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) print('copy {}'.format(name)) 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) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(args.gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=16, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=16, pin_memory=True) else: #sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = np.arange(train_dataset_size) np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) train_sampler_all = data.sampler.SubsetRandomSampler(train_ids) train_gt_sampler_all = data.sampler.SubsetRandomSampler(train_ids) train_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler( train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) trainloader_all = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler_all, num_workers=16, pin_memory=True) trainloader_gt_all = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler_all, num_workers=16, pin_memory=True) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=16, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=16, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=16, pin_memory=True) trainloader_remain_iter = iter(trainloader_remain) trainloader_all_iter = iter(trainloader_all) trainloader_iter = iter(trainloader) trainloader_gt_iter = iter(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 #y_real_, y_fake_ = Variable(torch.ones(args.batch_size, 1).cuda()), Variable(torch.zeros(args.batch_size, 1).cuda()) for i_iter in range(args.num_steps): loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_fm_value = 0 loss_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): # train G # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # train with source try: batch = next(trainloader_iter) except: trainloader_iter = iter(trainloader) batch = next(trainloader_iter) images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) #ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) loss_seg = loss_calc(pred, labels, args.gpu) loss_seg.backward() loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size if i_iter >= args.adv_start: #fm loss calc try: batch = next(trainloader_all_iter) except: trainloader_iter = iter(trainloader_all) batch = next(trainloader_all_iter) images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) #ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) _, D_out_y_pred = model_D(F.softmax(pred)) trainloader_gt_iter = iter(trainloader_gt) batch = next(trainloader_gt_iter) _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) #ignore_mask_gt = (labels_gt.numpy() == 255) _, D_out_y_gt = model_D(D_gt_v) fm_loss = torch.mean( torch.abs( torch.mean(D_out_y_gt, 0) - torch.mean(D_out_y_pred, 0))) loss = loss_seg + args.lambda_fm * fm_loss # proper normalization fm_loss.backward() #loss_seg_value += loss_seg.data.cpu().numpy()[0]/args.iter_size loss_fm_value += fm_loss.data.cpu().numpy()[0] / args.iter_size loss_value += loss.data.cpu().numpy()[0] / args.iter_size # train D # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred pred = pred.detach() D_out_z, _ = model_D(F.softmax(pred)) y_fake_ = Variable(torch.zeros(D_out_z.size(0), 1).cuda()) loss_D_fake = criterion(D_out_z, y_fake_) # train with gt # get gt labels _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) #ignore_mask_gt = (labels_gt.numpy() == 255) D_out_z_gt, _ = model_D(D_gt_v) #D_out = interp(D_out_x) y_real_ = Variable(torch.ones(D_out_z_gt.size(0), 1).cuda()) loss_D_real = criterion(D_out_z_gt, y_real_) loss_D = loss_D_fake + loss_D_real loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] optimizer.step() optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_D = {3:.3f}'. format(i_iter, args.num_steps, loss_seg_value, loss_D_value)) print('fm_loss: ', loss_fm_value, ' g_loss: ', loss_value) if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu print("Evaluating model") print(args.restore_from) print("classifier model") print(args.restore_from_classifier) print("sigmoid threshold") print(args.sigmoid_threshold) if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) model = Res_Deeplab(num_classes=args.num_classes) model_cls = Res_Deeplab_class(num_classes=args.num_classes, mode=6, latent_vars=args.latent_vars) saved_state_dict = torch.load(args.restore_from) model.load_state_dict(saved_state_dict, strict=False) model.eval() model.cuda(gpu0) saved_state_dict = torch.load(args.restore_from_classifier) model_cls.load_state_dict(saved_state_dict, strict=False) model_cls.eval() model_cls.cuda(gpu0) testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] colorize = VOCColorize() combo_matrix = np.zeros((args.num_classes, args.latent_vars + 1), dtype=np.float32) for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd' % (index)) image, label, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output).cpu().data[0].numpy() output = output[:, :size[0], :size[1]] cls_pred = F.sigmoid( model_cls(Variable(image, volatile=True).cuda(gpu0))) cls_pred = cls_pred.cpu().data.numpy()[0] gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int) #gt_classes = np.unique(gt).tolist() for clsID in range(1, args.num_classes): if cls_pred[clsID - 1] < args.sigmoid_threshold: output[clsID, :, :] = -1000000000 output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) filename = os.path.join(args.save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray( colorize(output).transpose(1, 2, 0), 'RGB') #color_file.save(filename) #filename = os.path.join(args.save_dir, '{}_lv.png'.format(name[0])) #color_file = Image.fromarray(colorize(output_lv).transpose(1, 2, 0), 'RGB') #color_file.save(filename) filename_gt = os.path.join(args.save_dir, '{}_gt.png'.format(name[0])) color_file_gt = Image.fromarray(colorize(gt).transpose(1, 2, 0), 'RGB') #color_file_gt.save(filename_gt) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) filename = os.path.join( args.save_dir, args.restore_from.split('/')[-1][:-4] + '_with_classifier_result.txt') confusion_matrix = get_iou(data_list, args.num_classes, filename)
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) model = Res_Deeplab(num_classes=args.num_classes) # if args.pretrained_model != None: # args.restore_from = pretrianed_models_dict[args.pretrained_model] # # 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 = Res_Deeplab(num_classes=args.num_classes) #model.load_state_dict(torch.load('/data/wyc/AdvSemiSeg/snapshots/VOC_15000.pth'))#70.7 state_dict = torch.load( '/data1/wyc/AdvSemiSeg/snapshots/VOC_t_baseline_1adv_mul_20000.pth' ) #baseline707 adv 709 nadv 705()*2#n adv0.694 # state_dict = torch.load( # '/home/wyc/VOC_t_baseline_nadv2_20000.pth') # baseline707 adv 709 nadv 705()*2 # original saved file with DataParallel # create new OrderedDict that does not contain `module.` from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v # load params new_params = model.state_dict().copy() for name, param in new_params.items(): print(name) if name in new_state_dict and param.size( ) == new_state_dict[name].size(): new_params[name].copy_(new_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.eval() model.cuda(gpu0) testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] colorize = VOCColorize() tag = 0 for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd' % (index)) image, label, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda(gpu0)) pred = interp(output) pred01 = F.softmax(pred, dim=1) output = interp(output).cpu().data[0].numpy() image = Variable(image).cuda() pred_re = F.softmax(pred, dim=1).repeat(1, 3, 1, 1) indices_1 = torch.index_select(image, 1, Variable(torch.LongTensor([0])).cuda()) indices_2 = torch.index_select(image, 1, Variable(torch.LongTensor([1])).cuda()) indices_3 = torch.index_select(image, 1, Variable(torch.LongTensor([2])).cuda()) img_re = torch.cat([ indices_1.repeat(1, 21, 1, 1), indices_2.repeat(1, 21, 1, 1), indices_3.repeat(1, 21, 1, 1), ], 1) mul_img = pred_re * img_re for i_l in range(label.shape[0]): label_set = np.unique(label[i_l]).tolist() for ls in label_set: if ls != 0 and ls != 255: ls = int(ls) img_p = torch.cat([ mul_img[i_l][ls].unsqueeze(0).unsqueeze(0), mul_img[i_l][ls + 21].unsqueeze(0).unsqueeze(0), mul_img[i_l][ls + 21 + 21].unsqueeze(0).unsqueeze(0) ], 1) imgs = img_p.squeeze() imgs = imgs.transpose(0, 1) imgs = imgs.transpose(1, 2) imgs = imgs.data.cpu().numpy() img_ori = image[0] img_ori = img_ori.squeeze() img_ori = img_ori.transpose(0, 1) img_ori = img_ori.transpose(1, 2) img_ori = img_ori.data.cpu().numpy() pred_ori = pred01[0][ls] pred_ori = pred_ori.data.cpu().numpy() pred_0 = pred_ori.copy() pred_ori = pred_ori size = pred_ori.shape color_image = np.zeros((3, size[0], size[1]), dtype=np.uint8) for i in range(size[0]): for j in range(size[1]): if pred_0[i][j] > 0.995: color_image[0][i][j] = 0 color_image[1][i][j] = 255 color_image[2][i][j] = 0 elif pred_0[i][j] > 0.9: color_image[0][i][j] = 255 color_image[1][i][j] = 0 color_image[2][i][j] = 0 elif pred_0[i][j] > 0.7: color_image[0][i][j] = 0 color_image[1][i][j] = 0 color_image[2][i][j] = 255 color_image = color_image.transpose((1, 2, 0)) # print pred_ori.shape cv2.imwrite( osp.join('/data1/wyc/AdvSemiSeg/vis/img_pred', name[0] + '.png'), imgs) cv2.imwrite( osp.join('/data1/wyc/AdvSemiSeg/vis/image', name[0] + '.png'), img_ori) cv2.imwrite( osp.join('/data1/wyc/AdvSemiSeg/vis/pred', name[0] + '.png'), color_image) output = output[:, :size[0], :size[1]] gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int) output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) filename = os.path.join(args.save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray( colorize(output).transpose(1, 2, 0), 'RGB') color_file.save(filename) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) filename = os.path.join(args.save_dir, 'result.txt') get_iou(data_list, args.num_classes, filename)
def main(): h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True gpu = args.gpu np.random.seed(args.random_seed) # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() for name, param in new_params.items(): print(name) if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.train() model.cuda(args.gpu) # init D model_D = FCDiscriminator(num_classes=args.num_classes) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(args.gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) # load dataset train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: #sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = np.arange(train_dataset_size) np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) # labeled data train_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) train_gt_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) # unlabeled data train_remain_sampler = data.sampler.SubsetRandomSampler( train_ids[partial_size:]) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() # loss/bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 for i_iter in range(args.num_steps): loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_D_ul_value = 0 loss_semi_value = 0 loss_semi_adv_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) # creating 2nd discriminator as a copy of the 1st one if i_iter == args.discr_split: model_D_ul = FCDiscriminator(num_classes=args.num_classes) model_D_ul.load_state_dict(net_D.state_dict()) model_D_ul.train() model_D_ul.cuda(args.gpu) optimizer_D_ul = optim.Adam(model_D_ul.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) # start training 2nd discriminator after specified number of steps if i_iter >= args.discr_split: optimizer_D_ul.zero_grad() adjust_learning_rate_D(optimizer_D_ul, i_iter) for sub_i in range(args.iter_size): # train Segmentation # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # don't accumulate grads in D_ul, in case split has already been made if i_iter >= args.discr_split: for param in model_D_ul.parameters(): param.requires_grad = False # do semi-supervised training first if args.lambda_semi_adv > 0 and i_iter >= args.semi_start_adv: try: _, batch = trainloader_remain_iter.next() except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = trainloader_remain_iter.next() # only access to img images, _, _, _ = batch images = Variable(images).cuda(args.gpu) pred = interp(model(images)) pred_remain = pred.detach() # choose discriminator depending on the iteration if i_iter >= args.discr_split: D_out = interp(model_D_ul(F.softmax(pred))) else: D_out = interp(model_D(F.softmax(pred))) D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze( axis=1) ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype( np.bool) # adversarial loss loss_semi_adv = args.lambda_semi_adv * bce_loss( D_out, make_D_label(gt_label, ignore_mask_remain, args.gpu)) loss_semi_adv = loss_semi_adv / args.iter_size # true loss value without multiplier loss_semi_adv_value += loss_semi_adv.data.cpu().numpy( ) / args.lambda_semi_adv loss_semi_adv.backward() else: loss_semi = None loss_semi_adv = None # train with labeled images try: _, batch = trainloader_iter.next() except: trainloader_iter = enumerate(trainloader) _, batch = trainloader_iter.next() images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) D_out = interp(model_D(F.softmax(pred))) # computing loss loss_seg = loss_calc(pred, labels, args.gpu) loss_adv_pred = bce_loss( D_out, make_D_label(gt_label, ignore_mask, args.gpu)) loss = loss_seg + args.lambda_adv_pred * loss_adv_pred # proper normalization loss = loss / args.iter_size loss.backward() loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size loss_adv_pred_value += loss_adv_pred.data.cpu().numpy( ) / args.iter_size # train D and D_ul # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True if i_iter >= args.discr_split: for param in model_D_ul.parameters(): param.requires_grad = True # train D with pred pred = pred.detach() # before split, traing D with both labeled and unlabeled if args.D_remain and i_iter < args.discr_split and ( args.lambda_semi > 0 or args.lambda_semi_adv > 0): pred = torch.cat((pred, pred_remain), 0) ignore_mask = np.concatenate((ignore_mask, ignore_mask_remain), axis=0) D_out = interp(model_D(F.softmax(pred))) loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask, args.gpu)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy() # train D_ul with pred on unlabeled if i_iter >= args.discr_split and (args.lambda_semi > 0 or args.lambda_semi_adv > 0): D_ul_out = interp(model_D_ul(F.softmax(pred_remain))) loss_D_ul = bce_loss( D_ul_out, make_D_label(pred_label, ignore_mask_remain, args.gpu)) loss_D_ul = loss_D_ul / args.iter_size / 2 loss_D_ul.backward() loss_D_ul_value += loss_D_ul.data.cpu().numpy() # get gt labels try: _, batch = trainloader_gt_iter.next() except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = trainloader_gt_iter.next() images_gt, labels_gt, _, _ = batch images_gt = Variable(images_gt).cuda(args.gpu) with torch.no_grad(): pred_l = interp(model(images_gt)) # train D with gt D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) ignore_mask_gt = (labels_gt.numpy() == 255) D_out = interp(model_D(D_gt_v)) loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt, args.gpu)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy() # train D_ul with pseudo_gt (gt are substituted for pred) if i_iter >= args.discr_split: D_ul_out = interp(model_D_ul(F.softmax(pred_l))) loss_D_ul = bce_loss( D_ul_out, make_D_label(gt_label, ignore_mask_gt, args.gpu)) loss_D_ul = loss_D_ul / args.iter_size / 2 loss_D_ul.backward() loss_D_ul_value += loss_D_ul.data.cpu().numpy() optimizer.step() optimizer_D.step() if i_iter >= args.discr_split: optimizer_D_ul.step() print('exp = {}'.format(args.snapshot_dir)) print( 'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_D_ul={5:.3f}, loss_semi = {6:.3f}, loss_semi_adv = {7:.3f}' .format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_D_ul_value, loss_semi_value, loss_semi_adv_value)) if i_iter >= args.num_steps - 1: print('save model ...') torch.save( net.state_dict(), osp.join( args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_' + str(args.random_seed) + '.pth')) torch.save( net_D.state_dict(), osp.join( args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_' + str(args.random_seed) + '_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') torch.save( net.state_dict(), osp.join( args.snapshot_dir, 'VOC_' + str(i_iter) + '_' + str(args.random_seed) + '.pth')) torch.save( net_D.state_dict(), osp.join( args.snapshot_dir, 'VOC_' + str(i_iter) + '_' + str(args.random_seed) + '_D.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): # 将参数的input_size 映射到整数,并赋值,从字符串转换到整数二元组 h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = False gpu = args.gpu # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) # 确保模型中参数的格式与要加载的参数相同 # 返回一个字典,保存着module的所有状态(state);parameters和persistent buffers都会包含在字典中,字典的key就是parameter和buffer的 names。 new_params = model.state_dict().copy() for name, param in new_params.items(): # print (name) if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) # print('copy {}'.format(name)) model.load_state_dict(new_params) # 设置为训练模式 model.train() cudnn.benchmark = True model.cuda(gpu) # init D model_D = FCDiscriminator(num_classes=args.num_classes) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: # sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = list(range(train_dataset_size)) # ? np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) # 写入文件 train_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler( train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # ??? # labels for adversarial training pred_label = 0 gt_label = 1 for i_iter in range(args.num_steps): print("Iter:", i_iter) loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_semi_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): # train G # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # do semi first if args.lambda_semi > 0 and i_iter >= args.semi_start: try: _, batch = next(trainloader_remain_iter) except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = next(trainloader_remain_iter) # only access to img images, _, _, _ = batch images = Variable(images).cuda(gpu) # images = Variable(images).cpu() pred = interp(model(images)) D_out = interp(model_D(F.softmax(pred))) D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze( axis=1) # produce ignore mask semi_ignore_mask = (D_out_sigmoid < args.mask_T) semi_gt = pred.data.cpu().numpy().argmax(axis=1) semi_gt[semi_ignore_mask] = 255 semi_ratio = 1.0 - float( semi_ignore_mask.sum()) / semi_ignore_mask.size print('semi ratio: {:.4f}'.format(semi_ratio)) if semi_ratio == 0.0: loss_semi_value += 0 else: semi_gt = torch.FloatTensor(semi_gt) loss_semi = args.lambda_semi * loss_calc( pred, semi_gt, args.gpu) loss_semi = loss_semi / args.iter_size loss_semi.backward() loss_semi_value += loss_semi.data.cpu().numpy( )[0] / args.lambda_semi else: loss_semi = None # train with source try: _, batch = next(trainloader_iter) except: trainloader_iter = enumerate(trainloader) _, batch = next(trainloader_iter) images, labels, _, _ = batch images = Variable(images).cuda(gpu) # images = Variable(images).cpu() ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) loss_seg = loss_calc(pred, labels, args.gpu) D_out = interp(model_D(F.softmax(pred))) loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask)) loss = loss_seg + args.lambda_adv_pred * loss_adv_pred # proper normalization loss = loss / args.iter_size loss.backward() loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size loss_adv_pred_value += loss_adv_pred.data.cpu().numpy( )[0] / args.iter_size # train D # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred pred = pred.detach() D_out = interp(model_D(F.softmax(pred))) loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] # train with gt # get gt labels try: _, batch = next(trainloader_gt_iter) except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = next(trainloader_gt_iter) _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) # D_gt_v = Variable(one_hot(labels_gt)).cpu() ignore_mask_gt = (labels_gt.numpy() == 255) D_out = interp(model_D(D_gt_v)) loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) loss_D = loss_D / args.iter_size / 2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] optimizer.step() optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) print( 'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}' .format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value)) if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_' + str(i_iter) + '_D.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True gpu = args.gpu np.random.seed(args.random_seed) # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() for name, param in new_params.items(): print(name) if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.train() model.cuda(args.gpu) cudnn.benchmark = True if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) # load dataset train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: #sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = np.arange(train_dataset_size) np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) # labeled data train_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) train_gt_sampler = data.sampler.SubsetRandomSampler( train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) # unlabeled data train_remain_sampler = data.sampler.SubsetRandomSampler( train_ids[partial_size:]) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # loss/bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') for i_iter in range(args.num_steps): loss_seg_value = 0 loss_unlabeled_seg_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) for sub_i in range(args.iter_size): # train Segmentation # train with labeled images try: _, batch = trainloader_iter.next() except: trainloader_iter = enumerate(trainloader) _, batch = trainloader_iter.next() images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) pred = interp(model(images)) # computing loss loss_seg = loss_calc(pred, labels, args.gpu) # proper normalization loss = loss_seg / args.iter_size loss.backward() loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size # train with unlabeled if args.lambda_semi > 0 and i_iter >= args.semi_start: try: _, batch = trainloader_remain_iter.next() except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = trainloader_remain_iter.next() # only access to img images, _, _, _ = batch images = Variable(images).cuda(args.gpu) pred = interp(model(images)) semi_gt = pred.data.cpu().numpy().argmax(axis=1) semi_gt = torch.FloatTensor(semi_gt) loss_unlabeled_seg = args.lambda_semi * loss_calc( pred, semi_gt, args.gpu) loss_unlabeled_seg = loss_unlabeled_seg / args.iter_size loss_unlabeled_seg.backward() loss_unlabeled_seg_value += loss_unlabeled_seg.data.cpu( ).numpy() / args.lambda_semi else: if args.lambda_semi > 0 and i_iter < args.semi_start: loss_unlabeled_seg_value = 0 else: loss_unlabeled_seg_value = None optimizer.step() print('exp = {}'.format(args.snapshot_dir)) print( 'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_unlabeled_seg = {3:.3f} ' .format(i_iter, args.num_steps, loss_seg_value, loss_unlabeled_seg_value)) if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join( args.snapshot_dir, 'VOC_' + str(args.num_steps) + '_' + str(args.lambda_semi) + '_' + str(args.random_seed) + '.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') torch.save( model.state_dict(), osp.join( args.snapshot_dir, 'VOC_' + str(i_iter) + '_' + str(args.lambda_semi) + '_' + str(args.random_seed) + '.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) model = Res_Deeplab(num_classes=args.num_classes) # if args.pretrained_model != None: # args.restore_from = pretrianed_models_dict[args.pretrained_model] # # 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 = Res_Deeplab(num_classes=args.num_classes) #model.load_state_dict(torch.load('/data/wyc/AdvSemiSeg/snapshots/VOC_15000.pth'))#70.7 state_dict = torch.load( '/data1/wyc/AdvSemiSeg/snapshots/VOC_t_baseline_1adv_mul_new_two_patch2_20000.pth' ) #baseline707 adv 709 nadv 705()*2#n adv0.694 # state_dict = torch.load( # '/home/wyc/VOC_t_baseline_nadv2_20000.pth') # baseline707 adv 709 nadv 705()*2 # original saved file with DataParallel # create new OrderedDict that does not contain `module.` from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v # load params new_params = model.state_dict().copy() for name, param in new_params.items(): print(name) if name in new_state_dict and param.size( ) == new_state_dict[name].size(): new_params[name].copy_(new_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.eval() model.cuda(gpu0) testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] colorize = VOCColorize() for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd' % (index)) image, label, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output).cpu().data[0].numpy() output = output[:, :size[0], :size[1]] gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int) output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) filename = os.path.join(args.save_dir, '{}.png'.format(name[0])) color_file = Image.fromarray( colorize(output).transpose(1, 2, 0), 'RGB') color_file.save(filename) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) filename = os.path.join(args.save_dir, 'result.txt') get_iou(data_list, args.num_classes, filename)
def main(): # LD build for summary # training_summary = tf.summary.FileWriter(os.path.join(SUMMARY_DIR, 'train')) # val_summary = tf.summary.FileWriter(os.path.join(SUMMARY_DIR, 'val')) # dice_placeholder = tf.placeholder(tf.float32, [], name='dice') # loss_placeholder = tf.placeholder(tf.float32, [], name='loss') # # image_placeholder = tf.placeholder(tf.float32, [400*2, 400*2], name='image') # # prediction_placeholder = tf.placeholder(tf.float32, [400*2, 400*2], name='prediction') # tf.summary.scalar('dice', dice_placeholder) # tf.summary.scalar('loss', loss_placeholder) # # tf.summary.image('image', image_placeholder, max_outputs=1) # # tf.summary.image('prediction', prediction_placeholder, max_outputs=1) # summary_op = tf.summary.merge_all() # config = tf.ConfigProto() # config.gpu_options.allow_growth = True # sess = tf.Session(config=config) perfix_name = 'Liver' h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True gpu = args.gpu # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() for name, param in new_params.items(): print(name) if name in saved_state_dict and param.size( ) == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.train() model.cuda(args.gpu) cudnn.benchmark = True # LD delete ''' # init D model_D = FCDiscriminator(num_classes=args.num_classes) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(args.gpu) ''' if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) # LD ADD start from dataset.LiverDataset.liver_dataset import LiverDataset user_name = 'ld' validation_interval = 800 max_steps = 1000000000 batch_size = 5 n_neighboringslices = 1 input_size = 400 output_size = 400 slice_type = 'axial' oversample = False # reset_counter = args.reset_counter label_of_interest = 1 label_required = 0 magic_number = 26.91 max_slice_tries_val = 0 max_slice_tries_train = 2 fuse_labels = True apply_crop = False train_data_dir = "/home/" + user_name + "/Documents/dataset/ISBI2017/media/nas/01_Datasets/CT/LITS/Training_Batch_2" test_data_dir = "/home/" + user_name + "/Documents/dataset/ISBI2017/media/nas/01_Datasets/CT/LITS/Training_Batch_1" train_dataset = LiverDataset(data_dir=train_data_dir, slice_type=slice_type, n_neighboringslices=n_neighboringslices, input_size=input_size, oversample=oversample, label_of_interest=label_of_interest, label_required=label_required, max_slice_tries=max_slice_tries_train, fuse_labels=fuse_labels, apply_crop=apply_crop, interval=validation_interval, is_training=True, batch_size=batch_size, data_augmentation=True) val_dataset = LiverDataset(data_dir=test_data_dir, slice_type=slice_type, n_neighboringslices=n_neighboringslices, input_size=input_size, oversample=oversample, label_of_interest=label_of_interest, label_required=label_required, max_slice_tries=max_slice_tries_val, fuse_labels=fuse_labels, apply_crop=apply_crop, interval=validation_interval, is_training=False, batch_size=batch_size) # LD ADD end # LD delete ''' train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: #sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = range(train_dataset_size) np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) ''' # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # LD delete ''' # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9,0.99)) optimizer_D.zero_grad() ''' # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size, input_size), mode='bilinear') if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(input_size, input_size), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size, input_size), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 loss_list = [] for i_iter in range(iter_start, args.num_steps): loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_semi_value = 0 loss_semi_adv_value = 0 num_prediction = 0 num_ground_truth = 0 num_intersection = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) # LD delete ''' optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) ''' for sub_i in range(args.iter_size): # train G # LD delete ''' # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # do semi first if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv : try: _, batch = trainloader_remain_iter.next() except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = trainloader_remain_iter.next() # only access to img images, _, _, _ = batch images = Variable(images).cuda(args.gpu) pred = interp(model(images)) pred_remain = pred.detach() D_out = interp(model_D(F.softmax(pred))) D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1) ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(np.bool) loss_semi_adv = args.lambda_semi_adv * bce_loss(D_out, make_D_label(gt_label, ignore_mask_remain)) loss_semi_adv = loss_semi_adv/args.iter_size #loss_semi_adv.backward() loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv if args.lambda_semi <= 0 or i_iter < args.semi_start: loss_semi_adv.backward() loss_semi_value = 0 else: # produce ignore mask semi_ignore_mask = (D_out_sigmoid < args.mask_T) semi_gt = pred.data.cpu().numpy().argmax(axis=1) semi_gt[semi_ignore_mask] = 255 semi_ratio = 1.0 - float(semi_ignore_mask.sum())/semi_ignore_mask.size print('semi ratio: {:.4f}'.format(semi_ratio)) if semi_ratio == 0.0: loss_semi_value += 0 else: semi_gt = torch.FloatTensor(semi_gt) loss_semi = args.lambda_semi * loss_calc(pred, semi_gt, args.gpu) loss_semi = loss_semi/args.iter_size loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi loss_semi += loss_semi_adv loss_semi.backward() else: loss_semi = None loss_semi_adv = None ''' # train with source # LD delete ''' try: _, batch = trainloader_iter.next() except: trainloader_iter = enumerate(trainloader) _, batch = trainloader_iter.next() images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) ''' batch_image, batch_label = train_dataset.get_next_batch() batch_image = np.transpose(batch_image, axes=(0, 3, 1, 2)) batch_image = np.concatenate( [batch_image, batch_image, batch_image], axis=1) # print('Shape: ', np.shape(batch_image)) batch_image_torch = torch.Tensor(batch_image) images = Variable(batch_image_torch).cuda(args.gpu) # LD delete # ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) pred_ny = pred.data.cpu().numpy() pred_ny = np.transpose(pred_ny, axes=(0, 2, 3, 1)) pred_label_ny = np.squeeze(np.argmax(pred_ny, axis=3)) # prepare for dice # print('Shape of gt is: ', np.shape(batch_label)) # print('Shape of pred is: ', np.shape(pred_ny)) # print('Shape of pred_label is: ', np.shape(pred_label_ny)) num_prediction += np.sum(np.asarray(pred_label_ny, np.uint8)) num_ground_truth += np.sum(np.asarray(batch_label >= 1, np.uint8)) num_intersection += np.sum( np.asarray( np.logical_and(batch_label >= 1, pred_label_ny >= 1), np.uint8)) # num_intersection += np.sum(np.asarray(batch_label >= 1, np.uint8) == np.asarray(pred_label_ny, np.uint8)) loss_seg = loss_calc(pred, batch_label, args.gpu) # LD delete ''' D_out = interp(model_D(F.softmax(pred))) loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask)) loss = loss_seg + args.lambda_adv_pred * loss_adv_pred ''' loss = loss_seg # print('Loss is: ', loss) # proper normalization loss = loss / args.iter_size loss.backward() # print('Loss of numpy is: ', loss_seg.data.cpu().numpy()) # print('Loss of numpy of zero is: ', loss_seg.data.cpu().numpy()) loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size loss_list.append(loss_seg_value) # loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0]/args.iter_size # train D # LD delete ''' # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred pred = pred.detach() if args.D_remain: pred = torch.cat((pred, pred_remain), 0) ignore_mask = np.concatenate((ignore_mask,ignore_mask_remain), axis = 0) D_out = interp(model_D(F.softmax(pred))) loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask)) loss_D = loss_D/args.iter_size/2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] # train with gt # get gt labels try: _, batch = trainloader_gt_iter.next() except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = trainloader_gt_iter.next() _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) ignore_mask_gt = (labels_gt.numpy() == 255) D_out = interp(model_D(D_gt_v)) loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) loss_D = loss_D/args.iter_size/2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] ''' optimizer.step() # optimizer_D.step() dice = (2 * num_intersection + 1e-7) / (num_prediction + num_ground_truth + 1e-7) print('exp = {}'.format(args.snapshot_dir)) print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}'.format( i_iter, args.num_steps, loss_seg_value)) print( 'dice: %.4f, num_prediction: %d, num_ground_truth: %d, num_intersection: %d' % (dice, num_prediction, num_ground_truth, num_intersection)) if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, perfix_name + str(args.num_steps) + '.pth')) # torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, perfix_name +str(args.num_steps)+'_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: print('taking snapshot ...') # torch.save(model.state_dict(), osp.join(args.snapshot_dir, perfix_name + str(i_iter)+'.pth')) save_model(model, args.snapshot_dir, perfix_name, i_iter, 2) # torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, perfix_name +str(i_iter)+'_D.pth')) # if i_iter % UPDATE_TENSORBOARD_INTERVAL and i_iter != 0: # # update tensorboard # feed_dict = { # dice_placeholder: dice, # loss_placeholder: np.mean(loss_list) # } # summery_value = sess.run(summary_op, feed_dict) # training_summary.add_summary(summery_value, i_iter) # training_summary.flush() # # # for validation # val_num_prediction = 0 # val_num_ground_truth = 0 # val_num_intersection = 0 # loss_list = [] # # for _ in range(VAL_EXECUTE_TIMES): # batch_image, batch_label = val_dataset.get_next_batch() # batch_image = np.transpose(batch_image, axes=(0, 3, 1, 2)) # batch_image = np.concatenate([batch_image, batch_image, batch_image], axis=1) # # print('Shape: ', np.shape(batch_image)) # batch_image_torch = torch.Tensor(batch_image) # images = Variable(batch_image_torch).cuda(args.gpu) # # # LD delete # # ignore_mask = (labels.numpy() == 255) # pred = interp(model(images)) # pred_ny = pred.data.cpu().numpy() # pred_ny = np.transpose(pred_ny, axes=(0, 2, 3, 1)) # pred_label_ny = np.squeeze(np.argmax(pred_ny, axis=3)) # val_num_prediction += np.sum(np.asarray(pred_label_ny, np.uint8)) # val_num_ground_truth += np.sum(np.asarray(batch_label >= 1, np.uint8)) # val_num_intersection += np.sum(np.asarray(np.logical_and(batch_label >= 1, pred_label_ny >= 1), np.uint8)) # # loss_seg = loss_calc(pred, batch_label, args.gpu) # loss_seg_value += loss_seg.data.cpu().numpy() / args.iter_size # loss_list.append(loss_seg) # dice = (2 * val_num_intersection + 1e-7) / (val_num_prediction + val_num_ground_truth + 1e-7) # feed_dict = { # dice_placeholder: dice, # loss_placeholder: np.mean(loss_list) # } # summery_value = sess.run(summary_op, feed_dict) # val_summary.add_summary(summery_value, i_iter) # val_summary.flush() # loss_list = [] training_summary.close() val_summary.close() end = timeit.default_timer() print(end - start, 'seconds')
def main(): """Create the model and start the training.""" w, h = map(int, args.input_size_source.split(',')) input_size_source = (w, h) w, h = map(int, args.input_size_target.split(',')) input_size_target = (w, h) cudnn.enabled = True # Create network if args.model == 'ResNet': 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) 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.load_state_dict(saved_state_dict) elif args.model == 'VGG': model = DeeplabVGG(num_classes=args.num_classes, pretrained=True, vgg16_caffe_path=args.restore_from) # saved_state_dict = torch.load(args.restore_from) # model.load_state_dict(saved_state_dict) optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() model.train() model.cuda(args.gpu) cudnn.benchmark = True #Discrimintator setting model_D = FCDiscriminator(num_classes=args.num_classes) model_D.train() model_D.cuda(args.gpu) optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() bce_loss = torch.nn.BCEWithLogitsLoss() # labels for adversarial training source_adv_label = 0 target_adv_label = 1 #Dataloader if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) trainloader = data.DataLoader(GTA5DataSet(args.translated_data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size, crop_size=input_size_source, 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) trainloader_iter = enumerate(trainloader) style_trainloader = data.DataLoader(GTA5DataSet( args.stylized_data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size, crop_size=input_size_source, 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) style_trainloader_iter = enumerate(style_trainloader) if STAGE == 1: targetloader = data.DataLoader(cityscapesDataSet( args.data_dir_target, args.data_list_target, max_iters=args.num_steps * args.iter_size * args.batch_size, crop_size=input_size_target, mean=IMG_MEAN, set=args.set), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) targetloader_iter = enumerate(targetloader) else: #Dataloader for self-training targetloader = data.DataLoader(cityscapesDataSetLabel( args.data_dir_target, args.data_list_target, max_iters=args.num_steps * args.iter_size * args.batch_size, crop_size=input_size_target, mean=IMG_MEAN, set=args.set, label_folder='Path to generated pseudo labels'), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) targetloader_iter = enumerate(targetloader) interp = nn.Upsample(size=(input_size_source[1], input_size_source[0]), mode='bilinear', align_corners=True) interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), mode='bilinear', align_corners=True) # load checkpoint model, model_D, optimizer, start_iter = load_checkpoint( model, model_D, optimizer, filename=args.snapshot_dir + 'checkpoint_' + CHECKPOINT + '.pth.tar') for i_iter in range(start_iter, args.num_steps): optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) #train segementation network # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # train with source if STAGE == 1: if i_iter % 2 == 0: _, batch = next(trainloader_iter) else: _, batch = next(style_trainloader_iter) else: _, batch = next(trainloader_iter) image_source, label, _, _ = batch image_source = Variable(image_source).cuda(args.gpu) pred_source = model(image_source) pred_source = interp(pred_source) loss_seg_source = loss_calc(pred_source, label, args.gpu) loss_seg_source_value = loss_seg_source.item() loss_seg_source.backward() if STAGE == 2: # train with target _, batch = next(targetloader_iter) image_target, target_label, _, _ = batch image_target = Variable(image_target).cuda(args.gpu) pred_target = model(image_target) pred_target = interp_target(pred_target) #target segmentation loss loss_seg_target = loss_calc(pred_target, target_label, gpu=args.gpu) loss_seg_target.backward() # optimize optimizer.step() if STAGE == 1: # train with target _, batch = next(targetloader_iter) image_target, _, _ = batch image_target = Variable(image_target).cuda(args.gpu) pred_target = model(image_target) pred_target = interp_target(pred_target) #output-level adversarial training D_output_target = model_D(F.softmax(pred_target)) loss_adv = bce_loss( D_output_target, Variable( torch.FloatTensor(D_output_target.data.size()).fill_( source_adv_label)).cuda(args.gpu)) loss_adv = loss_adv * args.lambda_adv loss_adv.backward() #train discriminator for param in model_D.parameters(): param.requires_grad = True pred_source = pred_source.detach() pred_target = pred_target.detach() D_output_source = model_D(F.softmax(pred_source)) D_output_target = model_D(F.softmax(pred_target)) loss_D_source = bce_loss( D_output_source, Variable( torch.FloatTensor(D_output_source.data.size()).fill_( source_adv_label)).cuda(args.gpu)) loss_D_target = bce_loss( D_output_target, Variable( torch.FloatTensor(D_output_target.data.size()).fill_( target_adv_label)).cuda(args.gpu)) loss_D_source = loss_D_source / 2 loss_D_target = loss_D_target / 2 loss_D_source.backward() loss_D_target.backward() #optimize optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) print('iter = {0:8d}/{1:8d}, loss_seg_source = {2:.5f}'.format( i_iter, args.num_steps, loss_seg_source_value)) if i_iter % args.save_pred_every == 0: print('taking snapshot ...') state = { 'iter': i_iter, 'model': model.state_dict(), 'model_D': model_D.state_dict(), 'optimizer': optimizer.state_dict() } torch.save( state, osp.join(args.snapshot_dir, 'checkpoint_' + str(i_iter) + '.pth.tar')) torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth')) torch.save( model_D.state_dict(), osp.join(args.snapshot_dir, 'GTA5_D_' + str(i_iter) + '.pth')) cityscapes_eval_dir = osp.join(args.cityscapes_eval_dir, str(i_iter)) if not os.path.exists(cityscapes_eval_dir): os.makedirs(cityscapes_eval_dir) eval(osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'), cityscapes_eval_dir, i_iter) iou19, iou13, iou = compute_mIoU(cityscapes_eval_dir, i_iter) outputfile = open(args.output_file, 'a') outputfile.write( str(i_iter) + '\t' + str(iou19) + '\t' + str(iou.replace('\n', ' ')) + '\n') outputfile.close()
def main(): h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True # create network model = Res_Deeplab(num_classes=args.num_classes) # load pretrained parameters 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) # only copy the params that exist in currendt model (caffe-like) new_params = model.state_dict().copy() for name, param in new_params.items(): print (name) if name in saved_state_dict and param.size() == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) print('copy {}'.format(name)) model.load_state_dict(new_params) model.train() model=nn.DataParallel(model) model.cuda() cudnn.benchmark = True # init D model_D = FCDiscriminator(num_classes=args.num_classes) if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D = nn.DataParallel(model_D) model_D.train() model_D.cuda() if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: #sample partial data partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print('loading train ids from {}'.format(args.partial_id)) else: train_ids = list(range(train_dataset_size)) np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.module.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum,weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9,0.99)) optimizer_D.zero_grad() # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 for i_iter in range(args.num_steps): loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_semi_value = 0 loss_semi_adv_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): # train G # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # do semi first if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv : try: _, batch = trainloader_remain_iter.next() except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = trainloader_remain_iter.next() # only access to img images, _, _, _ = batch images = Variable(images).cuda() pred = interp(model(images)) pred_remain = pred.detach() mask1=F.softmax(pred,dim=1).data.cpu().numpy() id2 = np.argmax(mask1, axis=1)#10, 321, 321) D_out = interp(model_D(F.softmax(pred,dim=1))) D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1) ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(np.bool) loss_semi_adv = args.lambda_semi_adv * bce_loss(D_out, make_D_label(gt_label, ignore_mask_remain)) loss_semi_adv = loss_semi_adv/args.iter_size #loss_semi_adv.backward() loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv if args.lambda_semi <= 0 or i_iter < args.semi_start: loss_semi_adv.backward() loss_semi_value = 0 else: # produce ignore mask semi_ignore_mask = (D_out_sigmoid < args.mask_T) #print semi_ignore_mask.shape 10,321,321 map2 = np.zeros([pred.size()[0], id2.shape[1], id2.shape[2]]) for k in range(pred.size()[0]): for i in range(id2.shape[1]): for j in range(id2.shape[2]): map2[k][i][j] = mask1[k][id2[k][i][j]][i][j] semi_ignore_mask = (map2 < 0.999999) semi_gt = pred.data.cpu().numpy().argmax(axis=1) semi_gt[semi_ignore_mask] = 255 semi_ratio = 1.0 - float(semi_ignore_mask.sum())/semi_ignore_mask.size print('semi ratio: {:.4f}'.format(semi_ratio)) if semi_ratio == 0.0: loss_semi_value += 0 else: semi_gt = torch.FloatTensor(semi_gt) loss_semi = args.lambda_semi * loss_calc(pred, semi_gt) loss_semi = loss_semi/args.iter_size loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi loss_semi += loss_semi_adv loss_semi.backward() else: loss_semi = None loss_semi_adv = None # train with source try: _, batch = trainloader_iter.next() except: trainloader_iter = enumerate(trainloader) _, batch = trainloader_iter.next() images, labels, _, _ = batch images = Variable(images).cuda() ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) loss_seg = loss_calc(pred, labels) D_out = interp(model_D(F.softmax(pred,dim=1))) loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask)) loss = loss_seg + args.lambda_adv_pred * loss_adv_pred # proper normalization loss = loss/args.iter_size loss.backward() loss_seg_value += loss_seg.data.cpu().numpy()[0]/args.iter_size loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0]/args.iter_size # train D # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred pred = pred.detach() if args.D_remain: pred = torch.cat((pred, pred_remain), 0) ignore_mask = np.concatenate((ignore_mask,ignore_mask_remain), axis = 0) D_out = interp(model_D(F.softmax(pred,dim=1))) loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask)) loss_D = loss_D/args.iter_size/2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] # train with gt # get gt labels try: _, batch = trainloader_gt_iter.next() except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = trainloader_gt_iter.next() _, labels_gt, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda() ignore_mask_gt = (labels_gt.numpy() == 255) D_out = interp(model_D(D_gt_v)) loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) loss_D = loss_D/args.iter_size/2 loss_D.backward() loss_D_value += loss_D.data.cpu().numpy()[0] optimizer.step() optimizer_D.step() print('exp = {}'.format(args.snapshot_dir)) print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}'.format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value, loss_semi_adv_value)) if i_iter >= args.num_steps-1: print( 'save model ...') torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+os.path.abspath(__file__).split('/')[-1].split('.')[0]+'_'+str(args.num_steps)+'.pth')) torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+os.path.abspath(__file__).split('/')[-1].split('.')[0]+'_'+str(args.num_steps)+'_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter!=0: print ('taking snapshot ...') torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+os.path.abspath(__file__).split('/')[-1].split('.')[0]+'_'+str(i_iter)+'.pth')) torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+os.path.abspath(__file__).split('/')[-1].split('.')[0]+'_'+str(i_iter)+'_D.pth')) end = timeit.default_timer() print(end-start,'seconds')
def main(): h, w = list(map(int, args.input_size.split(','))) # 321, 321 input_size = (h, w) cudnn.enabled = True # create network model = Res_Deeplab(num_classes=args.num_classes) # num_classes = 21 # load pretrained parameters 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) # only copy the params that exist in current model (caffe-like) new_params = model.state_dict().copy() for name, param in list(new_params.items()): # print(name) if name in saved_state_dict and param.size() == saved_state_dict[name].size(): new_params[name].copy_(saved_state_dict[name]) # print(('copy {}'.format(name))) 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) # num_classes = 21,全卷积判别模型 if args.restore_from_D is not None: model_D.load_state_dict(torch.load(args.restore_from_D)) model_D.train() model_D.cuda(args.gpu) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) if not os.path.exists('logs/'): os.makedirs('logs/') now_time = datetime.now().strftime('%Y-%m-%d-%H:%M:%S') log_file = 'logs/' + now_time + '.txt' file = open(log_file, 'w') # 保存loss train_dataset = VOCDataSet(args.data_dir, args.data_list, args.label_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) train_dataset_size = len(train_dataset) train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, args.label_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) if args.partial_data is None: # 使用全部数据 trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, # batch_size = 10 num_workers=5, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) else: # sample partial data 部分数据 partial_size = int(args.partial_data * train_dataset_size) if args.partial_id is not None: train_ids = pickle.load(open(args.partial_id)) print(('loading train ids from {}'.format(args.partial_id))) else: train_ids = list(range(train_dataset_size)) np.random.shuffle(train_ids) pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb')) # 将train_ids写入train_id.pkl train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:]) train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size]) trainloader = data.DataLoader(train_dataset, # 数据集中采样输入 batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True) trainloader_remain = data.DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True) trainloader_gt = data.DataLoader(train_gt_dataset, batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True) trainloader_remain_iter = enumerate(trainloader_remain) trainloader_iter = enumerate(trainloader) trainloader_gt_iter = enumerate(trainloader_gt) # implement model.optim_parameters(args) to handle different models' lr setting # optimizer for segmentation network optimizer = optim.SGD(model.optim_parameters(args), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() # optimizer for discriminator network optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99)) optimizer_D.zero_grad() # loss/ bilinear upsampling bce_loss = BCEWithLogitsLoss2d() # interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear') # labels for adversarial training pred_label = 0 gt_label = 1 best_loss = 1 best_epoch = 0 for i_iter in range(args.num_steps): # num_steps = 20000 loss_seg_value = 0 loss_adv_pred_value = 0 loss_D_value = 0 loss_semi_value = 0 loss_semi_adv_value = 0 optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) optimizer_D.zero_grad() adjust_learning_rate_D(optimizer_D, i_iter) for sub_i in range(args.iter_size): # iter_size = 1 # train G # don't accumulate grads in D for param in model_D.parameters(): param.requires_grad = False # do semi first if (args.lambda_semi > 0 or args.lambda_semi_adv > 0) and i_iter >= args.semi_start_adv: try: _, batch = next(trainloader_remain_iter) except: trainloader_remain_iter = enumerate(trainloader_remain) _, batch = next(trainloader_remain_iter) # only access to img 无标签数据 images, _, _, _, _ = batch images = Variable(images).cuda(args.gpu) pred = interp(model(images)) pred_remain = pred.detach() # 返回一个新的Variable,不具有grade D_out = interp(model_D(F.softmax(pred))) D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1) ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(np.bool) loss_semi_adv = args.lambda_semi_adv * bce_loss(D_out, make_D_label(gt_label, ignore_mask_remain)) loss_semi_adv = loss_semi_adv/args.iter_size # loss_semi_adv.backward() # print('bug,', loss_semi_adv.data.cpu().numpy()) # loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()/args.lambda_semi_adv if args.lambda_semi <= 0 or i_iter < args.semi_start: loss_semi_adv.backward() loss_semi_value = 0 else: # produce ignore mask semi_ignore_mask = (D_out_sigmoid < args.mask_T) # mask_T = 0.2,阈值 semi_gt = pred.data.cpu().numpy().argmax(axis=1) # 返回维度为1上的最大值的下标 semi_gt[semi_ignore_mask] = 255 semi_ratio = 1.0 - float(semi_ignore_mask.sum())/semi_ignore_mask.size # 被忽略的点占的比重 print(('semi ratio: {:.4f}'.format(semi_ratio))) if semi_ratio == 0.0: loss_semi_value += 0 else: semi_gt = torch.FloatTensor(semi_gt) loss_semi = args.lambda_semi * loss_calc(pred, semi_gt, args.gpu) loss_semi = loss_semi/args.iter_size # loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi loss_semi_value += loss_semi.data.cpu().numpy()/args.lambda_semi loss_semi += loss_semi_adv loss_semi.backward() else: loss_semi = None loss_semi_adv = None # train with source try: _, batch = next(trainloader_iter) except: trainloader_iter = enumerate(trainloader) _, batch = next(trainloader_iter) images, labels, _, _, _ = batch # 有标签数据 images = Variable(images).cuda(args.gpu) ignore_mask = (labels.numpy() == 255) pred = interp(model(images)) # interp上采样 loss_seg = loss_calc(pred, labels, args.gpu) # 语义分割的cross entropy loss # loss_seg_NLL = loss_NLL(pred, labels, args.gpu) # 语义分割的NLLLoss D_out = interp(model_D(F.softmax(pred))) # 得到判别模型输出的判别图 loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask)) loss = loss_seg + args.lambda_adv_pred * loss_adv_pred # proper normalization loss = loss/args.iter_size loss.backward() # loss_seg_value += loss_seg.data.cpu().numpy()[0]/args.iter_size # loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0]/args.iter_size loss_seg_value += loss_seg.data.cpu().numpy()/args.iter_size # loss_seg_value += loss_seg_NLL.data.cpu().numpy()/args.iter_size loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()/args.iter_size # train D # bring back requires_grad for param in model_D.parameters(): param.requires_grad = True # train with pred pred = pred.detach() if args.D_remain: pred = torch.cat((pred, pred_remain), 0) ignore_mask = np.concatenate((ignore_mask,ignore_mask_remain), axis=0) D_out = interp(model_D(F.softmax(pred))) loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask)) loss_D = loss_D/args.iter_size/2 loss_D.backward() # loss_D_value += loss_D.data.cpu().numpy()[0] loss_D_value += loss_D.data.cpu().numpy() # train with gt # get gt labels try: _, batch = next(trainloader_gt_iter) except: trainloader_gt_iter = enumerate(trainloader_gt) _, batch = next(trainloader_gt_iter) _, labels_gt, _, _, _ = batch D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu) # 每个类别一张label图,batch * class * h * w ignore_mask_gt = (labels_gt.numpy() == 255) D_out = interp(model_D(D_gt_v)) # ground_truth输入判别模型 loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt)) loss_D = loss_D/args.iter_size/2 loss_D.backward() # loss_D_value += loss_D.data.cpu().numpy()[0] loss_D_value += loss_D.data.cpu().numpy() optimizer.step() optimizer_D.step() print(('exp = {}'.format(args.snapshot_dir))) print(('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, ' 'loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}'.format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value, loss_semi_adv_value))) file.write('{0} {1} {2} {3} {4}\n'.format(loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value, loss_semi_adv_value)) if loss_seg_value < best_loss: # 保存最优模型,删除次优模型 # print('loss:', loss_seg_value, 'best:', best_loss) torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'VOC_epoch_{0}_seg_loss_{1}.pth'.format(i_iter+1, loss_seg_value))) torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_epoch_{0}_seg_loss_{1}_D.pth'.format(i_iter+1, loss_seg_value))) delete_models(best_epoch + 1, best_loss) best_loss = loss_seg_value best_epoch = i_iter if i_iter >= args.num_steps-1: # num_step = 20000 print('save model ...') torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'VOC_'+str(args.num_steps)+'.pth')) torch.save(model_D.state_dict(), osp.join(args.snapshot_dir, 'VOC_'+str(args.num_steps)+'_D.pth')) break if i_iter % args.save_pred_every == 0 and i_iter != 0: # save_pred_every = 5000 print('taking snapshot ...') torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(i_iter)+'.pth')) torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(i_iter)+'_D.pth')) end = timeit.default_timer() print(end-start, 'seconds') file.close()