def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) 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) interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd' % (index)) images, label, size, name = batch images = Variable(images, volatile=True) h, w, c = size[0].numpy() images075 = nn.Upsample(size=(int(h * 0.75), int(w * 0.75)), mode='bilinear')(images) images05 = nn.Upsample(size=(int(h * 0.5), int(w * 0.5)), mode='bilinear')(images) out100 = model(images.cuda(args.gpu)) out075 = model(images075.cuda(args.gpu)) out05 = model(images05.cuda(args.gpu)) o_h, o_w = out100.size()[2:] interpo1 = nn.Upsample(size=(o_h, o_w), mode='bilinear') out_max = torch.max(torch.stack( [out100, interpo1(out075), interpo1(out05)]), dim=0)[0] output = interp(out_max).cpu().data[0].numpy() output = output[:, :h, :w] output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) gt = np.asarray(label[0].numpy()[:h, :w], dtype=np.int) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) get_iou(data_list, args.num_classes)
def __init__(self, num_classes, train_sfe=True): super(SegFeatureExtraction, self).__init__() self.model = Res_Deeplab(num_classes=num_classes).layer4 if not train_sfe: # freeze parameters for param in self.model.parameters(): param.requires_grad = False
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from) model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) valid_loader = data.DataLoader(BerkeleyDataSet(args.data_dir, args.data_list, mean=IMG_MEAN, scale=False, mirror=False, train=False), batch_size=1, shuffle=False, pin_memory=True) test_loader = data.DataLoader(BerkeleyDataTestSet(args.data_dir, args.data_list, mean=IMG_MEAN)) interp = nn.Upsample(size=(321, 321), mode='bilinear', align_corners=True) data_list = [] # #Evaluation loop for Valid Loader with torch.no_grad(): for index, batch in enumerate(valid_loader): if index % 100 == 0: print('%d processd'%(index)) image, label, name, size = batch h, w, c = size[0].numpy() # print(name) output = model(Variable(image).cuda(gpu0)) output = interp(output).cpu().data[0].numpy() # print(output.shape) output = output.transpose(1,2,0) output = np.asarray(np.argmax(output, axis=2), dtype=np.int) # print(output) # print(label[0].shape) ground_truth = np.asarray(label[0].numpy()[:h,:w], dtype=np.int) show_all(ground_truth, output) data_list.append([ground_truth.flatten(), output.flatten()]) get_iou(data_list, args.num_classes) #Evaluation loop for Test Loader '''
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) 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) interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] with open('result.txt', 'w') as f: for index, batch in enumerate(testloader, 0): #print('%d processd'%(index)) image, label, size, name = batch size = size[0].numpy() #the size of original input image #print("size:",size) output = model(Variable(image, volatile=True).cuda(gpu0)) #print("model output size: ",output.size()) 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) prob = softmax(output, axis=2) entropy = (-prob * np.log(prob)).mean() * 100 output = np.asarray(np.argmax(output, axis=2), dtype=np.int) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) iou = get_iou(data_list, args.num_classes) data_list = [] print('{:.4f}, {:.4f}'.format(entropy.item(), iou.item())) f.write('{:.8f}, {:.8f}\n'.format(entropy.item(), iou.item()))
def main(): """Create the model and start the evaluation process.""" args = get_arguments() model = Res_Deeplab(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from) model.load_state_dict(saved_state_dict) model.eval() if torch.cuda.is_available(): args.device = torch.device('cuda') model = model.cuda() print() else: args.device = torch.device('cpu') print('the model is operating on {}'.format(args.device)) testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, hue_value=args.hue_value, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd' % (index)) image, label, size, name = batch image = image.to(device=args.device) size = size[0].numpy() output = model(Variable(image, volatile=True)) 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) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) get_iou(data_list, args.num_classes, args.hue_value)
def __init__(self, num_classes, train_cfe=True): super(CommonFeatureExtraction, self).__init__() self.model = Res_Deeplab(num_classes=num_classes) self.resnet_feature_layers = [ 'conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3' ] resnet_module_list = [ self.model.conv1, self.model.bn1, self.model.relu, self.model.maxpool, self.model.layer1, self.model.layer2, self.model.layer3 ] self.model = nn.Sequential(*resnet_module_list) if not train_cfe: # freeze parameters for param in self.model.parameters(): param.requires_grad = False
def iou(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) 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) interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True) data_list = [] for index, batch in enumerate(testloader): if index % 100 == 0: #print('%d processd'%(index)) sys.stdout.flush() image, label, size, name = batch size = size[0].numpy() with torch.no_grad(): output = model(Variable(image).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) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) return get_iou(data_list, args.num_classes)
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from) model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) dataset = GenericDataset(DATA_DIRECTORY, 'val', test_transform) testloader = data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True) interp = nn.Upsample(size=(505, 505), mode='bilinear') data_list = [] 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) # show_all(gt, output) data_list.append([gt.flatten(), output.flatten()]) get_iou(data_list, args.num_classes)
def main(): """Create the model and start the evaluation process.""" args = get_arguments() gpu0 = args.gpu model = Res_Deeplab(num_classes=args.num_classes) 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) interp = nn.Upsample(size=(505, 505), mode='bilinear') sigmoid = nn.Sigmoid() data_list = [] for index, batch in enumerate(testloader): if index % 1 == 0: print('%d processd' % (index)) image, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output) output = output[:, :size[0], :size[1]] output = torch.from_numpy(output).float() result = torchvision.transforms.ToPILImage()(output) path = os.path.join(SR_dir, name[0]) result.save(path)
def main(): """Create the model and start the evaluation process.""" gpu0 = 0 model = Res_Deeplab(num_classes=NUM_CLASSES) saved_state_dict = torch.load(RESTORE_FROM) model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) testloader = data.DataLoader(VOCDataSet(DATA_DIRECTORY, DATA_LIST_PATH, crop_size=(321, 321), mean=IMG_MEAN, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) interp = nn.Upsample(size=(321, 321), mode='bilinear', align_corners=True) #changed to model 321,321 data_list = [] for index, batch in enumerate(testloader): if index % 100 == 0: print('%d processd'%(index)) torch.save(batch, SAVE_TO + '/batch' + str(index) + '.pth') # Save the batch image, label, size, name = batch size = size[0].numpy() output = model(Variable(image, volatile=True).cuda(gpu0)) output = interp(output) torch.save(output, SAVE_TO + '/prediction' + str(index) + '.pth') #Save b11 prediction 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) data_list.append([gt.flatten(), output.flatten()]) get_iou(data_list, NUM_CLASSES)
else: os.mkdir(LR_dir) SR_dir = join(root_dir, 'SS_Output/SR') if os.path.isdir(SR_dir): pass else: os.mkdir(SR_dir) HR_dir = join(root_dir, 'SS_Output/HR') if os.path.isdir(HR_dir): pass else: os.mkdir(HR_dir) deeplab_res = Res_Deeplab(num_classes=21) saved_state_dict = torch.load('model/VOC12_scenes_20000.pth') deeplab_res.load_state_dict(saved_state_dict) deeplab_res = deeplab_res.eval() mid = mid_layer() criterion = CrossEntropy_Probability() criterion1 = CrossEntropy2d() criterion_kl = KL_Loss() mid = mid.cuda(gpuid) deeplab_res = deeplab_res.cuda(gpuid) criterion = criterion.cuda(gpuid) criterion1 = criterion1.cuda(gpuid) criterion_kl = criterion_kl.cuda(gpuid) testloader = data.DataLoader(VOCDataSSSet(root_dir,
def main(): """Create the model and start the training.""" 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) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. 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 == 21 or not i_parts[1] == 'layer5': new_params['.'.join(i_parts[1:])] = saved_state_dict[i] model.load_state_dict(new_params) #model.float() #model.eval() # use_global_stats = True model.train() model.cuda(args.gpu) cudnn.benchmark = True if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) trainloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.batch_size, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN), batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) optimizer = optim.SGD([{ 'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }, { 'params': get_10x_lr_params(model), 'lr': 10 * args.learning_rate }], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() interp = nn.Upsample(size=input_size, mode='bilinear') for i_iter, batch in enumerate(trainloader): images, labels, _, _ = batch images = Variable(images).cuda(args.gpu) optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) pred = interp(model(images)) loss = loss_calc(pred, labels, args.gpu) loss.backward() optimizer.step() print 'iter = ', i_iter, 'of', args.num_steps, 'completed, loss = ', loss.data.cpu( ).numpy() if i_iter >= args.num_steps - 1: print 'save model ...' torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC12_scenes_' + str(args.num_steps) + '.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, 'VOC12_scenes_' + str(i_iter) + '.pth')) end = timeit.default_timer() print end - start, 'seconds'
def main(): """Create the model and start the training.""" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True torch.manual_seed(args.random_seed) model = Res_Deeplab(num_classes=args.num_classes) model = torch.nn.DataParallel(model) optimizer = optim.SGD([{ 'params': get_1x_lr_params_NOscale(model.module), 'lr': args.learning_rate }, { 'params': get_10x_lr_params(model.module), 'lr': 10 * args.learning_rate }], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) if args.fine_tune: # fine tune from coco dataset 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('.') if i_parts[1] != 'layer5': new_params[i.replace('Scale', 'module')] = saved_state_dict[i] model.load_state_dict(new_params) elif args.restore_from: if os.path.isfile(args.restore_from): print("=> loading checkpoint '{}'".format(args.restore_from)) checkpoint = torch.load(args.restore_from) try: if args.set_start: args.start_step = int( math.ceil(checkpoint['example'] / args.batch_size)) model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}' (step {})".format( args.restore_from, args.start_step)) except: model.load_state_dict(checkpoint) print("=> loaded checkpoint '{}'".format(args.restore_from)) else: print("=> no checkpoint found at '{}'".format(args.restore_from)) if not args.is_training: # Frozen BN # when training, the model will use the running means and the # running vars of the pretrained model. # But note that eval() doesn't turn off history tracking. print( "Freezing BN layers, that is taking BN as linear transform layer") model.eval() else: model.train() model.cuda() cudnn.benchmark = True trainloader = data.DataLoader(XiangyaTrain(args.data_list, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, color_jitter=args.random_jitter, rotate=args.random_rotate, mean=IMG_MEAN), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() pixel_losses = AverageMeter() patch_losses = AverageMeter() accuracy = AverageMeter() writer = SummaryWriter(args.snapshot_dir) cnt = 0 actual_step = args.start_step while actual_step < args.final_step: iter_end = timeit.default_timer() for i_iter, (images, labels, patch_name) in enumerate(trainloader): actual_step = int(args.start_step + cnt) data_time.update(timeit.default_timer() - iter_end) images = Variable(images).cuda() optimizer.zero_grad() adjust_learning_rate(optimizer, actual_step) # pred = interp(model(images)) pred = model(images) image = images.data.cpu().numpy()[0] del images # 0 Normal 1 DG 2 JR labels = resize_target(labels, pred.size(2)) pixel_loss, patch_loss = loss_calc(pred, labels) loss = pixel_loss.double() + args.loss_coeff * patch_loss losses.update(loss.item(), pred.size(0)) pixel_losses.update(pixel_loss.item(), pred.size(0)) patch_losses.update(patch_loss.item(), pred.size(0)) acc = _pixel_accuracy(pred.data.cpu().numpy(), labels.data.cpu().numpy()) accuracy.update(acc, pred.size(0)) loss.backward() optimizer.step() batch_time.update(timeit.default_timer() - iter_end) iter_end = timeit.default_timer() if actual_step % args.print_freq == 0: print( 'iter: [{0}]{1}/{2}\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'Data {data_time.val:.3f} ({data_time.avg:.3f})\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'Pixel Loss {pixel_loss.val:.4f} ({pixel_loss.avg:.4f})\t' 'Patch_Loss {patch_loss.val:.4f} ({patch_loss.avg:.4f})\t' 'Pixel Accuracy {accuracy.val:.3f} ({accuracy.avg:.3f})'. format(cnt, actual_step, args.final_step, batch_time=batch_time, data_time=data_time, loss=losses, pixel_loss=pixel_losses, patch_loss=patch_losses, accuracy=accuracy)) writer.add_scalar("train_loss", losses.avg, actual_step) writer.add_scalar("pixel_loss", pixel_losses.avg, actual_step) writer.add_scalar("patch_loss", patch_losses.avg, actual_step) writer.add_scalar("pixel_accuracy", accuracy.avg, actual_step) writer.add_scalar("lr", optimizer.param_groups[0]['lr'], actual_step) # TODO complete this part using writer if actual_step % args.save_img_freq == 0: msk_size = pred.size(2) image = image.transpose(1, 2, 0) image = cv2.resize(image, (msk_size, msk_size), interpolation=cv2.INTER_NEAREST) image += IMG_MEAN image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) label = labels.data.cpu().numpy()[0] label = np.repeat(254, msk_size) - label * 127 single_pred = pred.data.cpu().numpy()[0].argmax(axis=0) single_pred = single_pred * 127 new_im = Image.new('RGB', (msk_size * 3, msk_size)) new_im.paste(Image.fromarray(image.astype('uint8'), 'RGB'), (0, 0)) new_im.paste(Image.fromarray(single_pred.astype('uint8'), 'L'), (msk_size, 0)) new_im.paste(Image.fromarray(label.astype('uint8'), 'L'), (msk_size * 2, 0)) new_im_name = 'B' + str(args.batch_size) + '_S' + str( actual_step) + '_' + patch_name[0] new_im_file = os.path.join(args.img_dir, new_im_name) new_im.save(new_im_file) if actual_step % args.save_pred_every == 0 and cnt != 0: print('taking snapshot ...') torch.save( { 'example': actual_step * args.batch_size, 'state_dict': model.state_dict() }, osp.join( args.snapshot_dir, 'Xiangya_Deeplab_B' + str(args.batch_size) + '_S' + str(actual_step) + '.pth')) cnt += 1 print('save the final model ...') torch.save( { 'example': actual_step * args.batch_size, 'state_dict': model.state_dict() }, osp.join( args.snapshot_dir, 'Xiangya_Deeplab_B' + str(args.batch_size) + '_S' + str(actual_step) + '.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): """Create the model and start the evaluation process.""" device = torch.device("cuda:" + str(args.gpu)) if not os.path.exists(args.save): os.makedirs(args.save) logger = util.set_logger(args.save, args.log_file, args.debug) logger.info('start with arguments %s', args) x_num = 0 with open(args.data_list) as f: for _ in f.readlines(): x_num = x_num + 1 sys.path.insert(0, 'dataset/helpers') if args.data_src == 'gta' or args.data_src == 'cityscapes': from labels import id2label, trainId2label elif args.data_src == 'synthia': from labels_cityscapes_synthia import id2label, trainId2label # label_2_id = 255 * np.ones((256, )) for l in id2label: if l in (-1, 255): continue label_2_id[l] = id2label[l].trainId id_2_label = np.array( [trainId2label[_].id for _ in trainId2label if _ not in (-1, 255)]) valid_labels = sorted(set(id_2_label.ravel())) scorer = ScoreUpdater(valid_labels, args.num_classes, x_num, logger) scorer.reset() if args.model == 'DeeplabRes': model = Res_Deeplab(num_classes=args.num_classes) # elif args.model == 'DeeplabVGG': # model = DeeplabVGG(num_classes=args.num_classes) # if args.restore_from == RESTORE_FROM: # args.restore_from = RESTORE_FROM_VGG if args.restore_from[:4] == 'http': saved_state_dict = model_zoo.load_url(args.restore_from) new_params = model.state_dict().copy() for i in saved_state_dict: # Scale.layer5.conv2d_list.3.weight i_parts = str(i).split('.') # print i_parts if not i_parts[0] == 'fc': new_params['.'.join(i_parts[0:])] = saved_state_dict[i] else: loc = "cuda:" + str(args.gpu) saved_state_dict = torch.load(args.restore_from, map_location=loc) new_params = saved_state_dict.copy() model.load_state_dict(new_params) #model.train() model.eval() model.to(device) testloader = data.DataLoader(GTA5TestDataSet(args.data_dir, args.data_list, test_scale=1.0, test_size=(1024, 512), mean=IMG_MEAN, std=IMG_STD, scale=False, mirror=False), batch_size=1, shuffle=False, pin_memory=True) test_scales = [float(_) for _ in str(args.test_scale).split(',')] h, w = map(int, args.test_image_size.split(',')) if version.parse(torch.__version__) >= version.parse('0.4.0'): interp = nn.Upsample(size=(h, w), mode='bilinear', align_corners=True) else: interp = nn.Upsample(size=(h, w), mode='bilinear') test_image_size = (h, w) mean_rgb = IMG_MEAN[::-1].copy() std_rgb = IMG_STD[::-1].copy() with torch.no_grad(): for index, batch in enumerate(testloader): image, label, _, name = batch img = image.clone() num_scales = len(test_scales) # output_dict = {k: [] for k in range(num_scales)} for scale_idx in range(num_scales): if version.parse(torch.__version__) > version.parse('0.4.0'): image = F.interpolate(image, scale_factor=test_scales[scale_idx], mode='bilinear', align_corners=True) else: test_size = (int(h * test_scales[scale_idx]), int(w * test_scales[scale_idx])) interp_tmp = nn.Upsample(size=test_size, mode='bilinear', align_corners=True) image = interp_tmp(img) if args.model == 'DeeplabRes': output2 = model(image.to(device)) coutput = interp(output2).cpu().data[0].numpy() if args.test_flipping: output2 = model( torch.from_numpy( image.numpy()[:, :, :, ::-1].copy()).to(device)) coutput = 0.5 * ( coutput + interp(output2).cpu().data[0].numpy()[:, :, ::-1]) if scale_idx == 0: output = coutput.copy() else: output += coutput output = output / num_scales output = output.transpose(1, 2, 0) output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8) pred_label = output.copy() label = label_2_id[np.asarray(label.numpy(), dtype=np.uint8)] scorer.update(pred_label.flatten(), label.flatten(), index) 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]))
# Prepare it for the model img_h, img_w, _ = image.shape inter_size = (img_h, img_w) image = image.transpose((2, 0, 1)) image = torch.from_numpy(image) # Change the single picture to a batch (otherwise a Pytorch Net cannot work with it) image = image.unsqueeze(0) # Prepare for the use of the GPU gpu0 = 0 os.environ["CUDA_VISIBLE_DEVICES"]=str(gpu0) cudnn.enabled = True # Initilize model model = Res_Deeplab(num_classes=21) # We use the path to the state dictionary on our own server as a default, can be changed with an argument pathToTrainedModel = '/root/20000StepsDefaultParametersBatch6/VOC12_scenes_20000.pth' if args.PathToPretrainedModel: pathToTrainedModel = PathToPretrainedModel saved_state_dict = torch.load(pathToTrainedModel) model.load_state_dict(saved_state_dict) #prepare for evaluation and use of cuda (GPU) model.eval() model.cuda() image = Variable(image).cuda() #gets and saves a gpu output interp = nn.Upsample(size=inter_size, mode='bilinear', align_corners=True)
def main(): print( "SRResNet with Semantic_KL_loss training from scratch on VOC 160*160 patches." ) #改参数之前改一下 global opt, model, netContent, deeplab_res, mid, semantic_criterion, semantic_kl_criterion, KL_DivLoss opt = parser.parse_args() print(opt) gpuid = 0 cuda = True opt.vgg_loss = False opt.semantic_loss = True KL_DivLoss = True if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") opt.seed = random.randint(1, 10000) print("Random Seed: ", opt.seed) torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) cudnn.benchmark = True if opt.vgg_loss: print('===> Loading VGG model') netVGG = models.vgg19() netVGG.load_state_dict( model_zoo.load_url( 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth') ) #建立vgg网络 从哪取loss照这个vgg loss写 class _content_model(nn.Module): def __init__(self): super(_content_model, self).__init__() self.feature = nn.Sequential( *list(netVGG.features.children())[:-1]) def forward(self, x): out = self.feature(x) return out netContent = _content_model() if opt.semantic_loss: print('===> Loading DeepLab model') deeplab_res = Res_Deeplab(num_classes=21) saved_state_dict = torch.load('model/VOC12_scenes_20000.pth') deeplab_res.load_state_dict(saved_state_dict) deeplab_res = deeplab_res.eval() semantic_criterion = CrossEntropy_Probability() semantic_kl_criterion = nn.KLDivLoss(size_average=False) print("===> Building model") model = Net() criterion = nn.MSELoss(size_average=False) print("===> Setting GPU") if cuda: model = model.cuda(gpuid) criterion = criterion.cuda(gpuid) if opt.vgg_loss: netContent = netContent.cuda(gpuid) if opt.semantic_loss: deeplab_res = deeplab_res.cuda(gpuid) semantic_criterion = semantic_criterion.cuda(gpuid) semantic_kl_criterion = semantic_kl_criterion.cuda(gpuid) # optionally resume from a checkpoint if opt.resume: if os.path.isfile(opt.resume): print("=> loading checkpoint '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) opt.start_epoch = checkpoint["epoch"] + 1 model.load_state_dict(checkpoint["model"].state_dict()) else: print("=> no checkpoint found at '{}'".format(opt.resume)) # optionally copy weights from a checkpoint if opt.pretrained: if os.path.isfile(opt.pretrained): print("=> loading model '{}'".format(opt.pretrained)) weights = torch.load(opt.pretrained) model.load_state_dict(weights['model'].state_dict()) else: print("=> no model found at '{}'".format(opt.pretrained)) print("===> Setting Optimizer") optimizer = optim.Adam(model.parameters(), lr=opt.lr) print("===> Training1") #root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5' root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train_hrlabel160_HDF5' files_num = len(os.listdir(root_dir)) for epoch in range(opt.start_epoch, opt.nEpochs + 1): #save_checkpoint(model, epoch) print("===> Loading datasets") x = random.sample(os.listdir(root_dir), files_num) for index in range(0, files_num): train_path = os.path.join(root_dir, x[index]) print("===> Training datasets: '{}'".format(train_path)) train_set = DatasetFromHdf5(train_path) #看 training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) avgloss = train(training_data_loader, optimizer, model, criterion, epoch, gpuid) if epoch % 2 == 0: save_checkpoint(model, epoch)
def main(): """Create the model and start the training.""" 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) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. 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 == 21 or not i_parts[1] == 'layer5': new_params['.'.join(i_parts[1:])] = saved_state_dict[i] model.load_state_dict(new_params) #model.float() #model.eval() # use_global_stats = True model.train() model.cuda(args.gpu) cudnn.benchmark = True if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) trainloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN), batch_size=args.batch_size, shuffle=True, num_workers=1, pin_memory=True) optimizer = optim.SGD([{ 'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }, { 'params': get_10x_lr_params(model), 'lr': 10 * args.learning_rate }], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() b_loss = 0 for i_iter, batch in enumerate(trainloader): images, labels, _, _ = batch images, labels = Variable(images), labels.numpy() h, w = images.size()[2:] images075 = nn.Upsample(size=(int(h * 0.75), int(w * 0.75)), mode='bilinear')(images) images05 = nn.Upsample(size=(int(h * 0.5), int(w * 0.5)), mode='bilinear')(images) out = model(images.cuda(args.gpu)) out075 = model(images075.cuda(args.gpu)) out05 = model(images05.cuda(args.gpu)) o_h, o_w = out.size()[2:] interpo1 = nn.Upsample(size=(o_h, o_w), mode='bilinear') interpo2 = nn.Upsample(size=(h, w), mode='bilinear') out_max = interpo2( torch.max(torch.stack([out, interpo1(out075), interpo1(out05)]), dim=0)[0]) loss = loss_calc(out_max, labels, args.gpu) d1, d2 = float(labels.shape[1]), float(labels.shape[2]) loss100 = loss_calc( out, nd.zoom(labels, (1.0, out.size()[2] / d1, out.size()[3] / d2), order=0), args.gpu) loss075 = loss_calc( out075, nd.zoom(labels, (1.0, out075.size()[2] / d1, out075.size()[3] / d2), order=0), args.gpu) loss05 = loss_calc( out05, nd.zoom(labels, (1.0, out05.size()[2] / d1, out05.size()[3] / d2), order=0), args.gpu) loss_all = (loss + loss100 + loss075 + loss05) / args.iter_size loss_all.backward() b_loss += loss_all.data.cpu().numpy() b_iter = i_iter / args.iter_size if b_iter >= args.num_steps - 1: print 'save model ...' optimizer.step() torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC12_scenes_' + str(args.num_steps) + '.pth')) break if i_iter % args.iter_size == 0 and i_iter != 0: print 'iter = ', b_iter, 'of', args.num_steps, 'completed, loss = ', b_loss optimizer.step() adjust_learning_rate(optimizer, b_iter) optimizer.zero_grad() b_loss = 0 if i_iter % (args.save_pred_every * args.iter_size) == 0 and b_iter != 0: print 'taking snapshot ...' torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC12_scenes_' + str(b_iter) + '.pth')) end = timeit.timeit() print end - start, 'seconds'
def main(): """Create the model and start the training.""" os.environ["CUDA_VISIBLE_DEVICES"]=str(args.gpu) 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) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Load pretrained COCO dataset saved_state_dict = torch.load(args.restore_from) new_params = model.state_dict().copy() # for param_tensor in model.state_dict(): # print(model.state_dict()[param_tensor].size()) for i in saved_state_dict: #Scale.layer5.conv2d_list.3.weight i_parts = i.split('.') # print i_parts # layer 5 is end of network if not args.num_classes == 3 or not i_parts[1]=='layer5': new_params['.'.join(i_parts[1:])] = saved_state_dict[i] model.load_state_dict(new_params) model.cuda() # Moves all parameters into GPU cudnn.benchmark = True # Creating folder for saved parameters of model if not os.path.exists(args.snapshot_dir): print("Creating Checkpoint Folder") os.makedirs(args.snapshot_dir) berkeleyDataset = BerkeleyDataSet(args.data_dir, args.data_list, max_iters=args.num_steps*args.batch_size, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN) trainloader = data.DataLoader(berkeleyDataset, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) optimizer = optim.SGD([{'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }, {'params': get_10x_lr_params(model), 'lr': 10*args.learning_rate}], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() interp = nn.Upsample(size=input_size, mode='bilinear') model.train() train_loss = 0 for i_iter, batch in enumerate(trainloader): images, labels, _, _ = batch images = Variable(images).cuda() optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) pred = interp(model(images)) loss = loss_calc(pred, labels) loss.backward() optimizer.step() print('iter = ', i_iter, 'of', args.num_steps,'completed, loss = ', loss.item()*images.size(0)) if i_iter >= args.num_steps-1: print('save model ...') torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'BDD_'+str(args.num_steps)+'.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, 'BDD_'+str(i_iter)+'.pth')) end = timeit.default_timer() print(end-start,'seconds')
def main(): """Create the model and start the training.""" if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) logger = util.set_logger(args.snapshot_dir, args.log_file, args.debug) logger.info('start with arguments %s', args) h, w = map(int, args.input_size.split(',')) input_size = (h, w) lscale, hscale = map(float, args.train_scale.split(',')) train_scale = (lscale, hscale) cudnn.enabled = True # Create network. model = Res_Deeplab(num_classes=args.num_classes) #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 == 21 or not i_parts[1]=='layer5': # new_params['.'.join(i_parts[1:])] = saved_state_dict[i] model_urls = { 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', } saved_state_dict = torch.utils.model_zoo.load_url(model_urls['resnet101']) # coco pretrained parameters: # 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 = str(i).split('.') # print i_parts if not i_parts[0] == 'fc': new_params['.'.join(i_parts[0:])] = saved_state_dict[i] model.load_state_dict(new_params) #model.float() model.eval() # use_global_stats = True #model.train() device = torch.device("cuda:" + str(args.gpu)) model.to(device) cudnn.benchmark = True trainloader = data.DataLoader(GTA5DataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.batch_size, crop_size=input_size, train_scale=train_scale, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN, std=IMG_STD), batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=args.pin_memory) optimizer = optim.SGD([{ 'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }, { 'params': get_10x_lr_params(model), 'lr': 10 * args.learning_rate }], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() interp = nn.Upsample(size=input_size, mode='bilinear', align_corners=True) for i_iter, batch in enumerate(trainloader): images, labels, _, _ = batch images = images.to(device) labels = labels.long().to(device) optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) pred = interp(model(images)) loss = loss_calc(pred, labels) loss.backward() optimizer.step() # print('iter = ', i_iter, 'of', args.num_steps,'completed, loss = ', loss.data.cpu().numpy()) logger.info('iter = {} of {} completed, loss = {:.4f}'.format( i_iter, args.num_steps, loss.data.cpu().numpy())) if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC12_scenes_' + str(args.num_steps) + '.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, 'VOC12_scenes_' + str(i_iter) + '.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 model = Res_Deeplab(num_classes=args.num_classes) saved_state_dict = torch.load(args.restore_from) model.load_state_dict(saved_state_dict) model.eval() model.cuda(gpu0) interp = nn.Upsample(size=(720, 1280), mode='bilinear', align_corners=True) data_list = [] ''' Use the following Data Directories and List Path for validation (BDD Dataset) DATA_DIRECTORY = './BDD_Deepdrive/bdd100k/' DATA_LIST_PATH = './dataset/list/BDD_val.txt' ''' # valid_loader = data.DataLoader(BerkeleyDataSet(args.data_dir, args.data_list, mean=IMG_MEAN, scale=False, mirror=False, train=False), # batch_size=1, shuffle=False, pin_memory=True) #Evaluation loop for Valid Loader # with torch.no_grad(): # for index, batch in enumerate(valid_loader): # if index % 100 == 0: # print('%d processd'%(index)) # image, label, name, size = batch # h, w, c = size[0].numpy() # # print(name) # output = model(Variable(image).cuda(gpu0)) # output = interp(output).cpu().data[0].numpy() # # print(output.shape) # output = output.transpose(1,2,0) # output = np.asarray(np.argmax(output, axis=2), dtype=np.int) # ground_truth = np.asarray(label[0].numpy()[:h,:w], dtype=np.int) # show_all(ground_truth, output, name) # data_list.append([ground_truth.flatten(), output.flatten()]) # get_iou(data_list, args.num_classes) ''' Use the following codes on the testset. Use the following Data Directories and List Path for testing (BerkeleyDataTestSet) DATA_DIRECTORY = './BDD_Deepdrive/bdd100k/' DATA_LIST_PATH = './dataset/list/BDD_test.txt' test_loader = data.DataLoader(BerkeleyDataTestSet(args.data_dir, args.data_list, mean=IMG_MEAN)) Use the following for Cityscapes Dataset (CityscapesDataSet) DATA_DIRECTORY = './CityScapes/stuttgart_01/' DATA_LIST_PATH = './dataset/list/CityScapesStuttgart_01.txt' test_loader = data.DataLoader(CityscapesDataSet(args.data_dir, args.data_list, mean=IMG_MEAN)) ''' test_loader = data.DataLoader( BerkeleyDataTestSet(args.data_dir, args.data_list, mean=IMG_MEAN)) masked_directory = 'D:/PyTorch-DeepLab-Berkeley/BDDMasked/' # #Evaluation loop for Test Loader with torch.no_grad(): for index, batch in enumerate(test_loader): if index % 100 == 0: print('%d processd' % (index)) image, name, size = batch h, w, c = size[0].numpy() # print(name) 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.int) # show_all(output, name=name) overlay( name=name, original_directory='./BDD_Deepdrive/bdd100k/images/100k/test/', masked_directory=masked_directory)
def main(): """Create the model and start the training.""" print("=====> Set GPU for training") if args.cuda: print("====> use gpu id: '{}'".format(args.gpus)) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus if not torch.cuda.is_available(): raise Exception( "No GPU found or Wrong gpu id, please run without --cuda") args.seed = random.randint(1, 10000) print("Random Seed: ", args.seed) torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) h, w = map(int, args.input_size.split(',')) input_size = (h, w) cudnn.enabled = True # Create network. print("=====> Building network") model = Res_Deeplab(num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. print("=====> Loading pretrained weights") 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 == 21 or not i_parts[1] == 'layer5': new_params['.'.join(i_parts[1:])] = saved_state_dict[i] model.load_state_dict(new_params) model.train() model.cuda() cudnn.benchmark = True if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) print('====> Computing network parameters') total_paramters = netParams(model) print('Total network parameters: ' + str(total_paramters)) print("=====> Preparing training data") trainloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.batch_size, crop_size=input_size, scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN), batch_size=args.batch_size, shuffle=True, num_workers=2, pin_memory=True) optimizer = optim.SGD([{ 'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }, { 'params': get_10x_lr_params(model), 'lr': 10 * args.learning_rate }], lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() interp = nn.Upsample(size=input_size, mode='bilinear') logFileLoc = args.snapshot_dir + args.logFile if os.path.isfile(logFileLoc): logger = open(logFileLoc, 'a') else: logger = open(logFileLoc, 'w') logger.write("Parameters: %s" % (str(total_paramters))) logger.write("\n%s\t\t%s" % ('iter', 'Loss(train)')) logger.flush() print("=====> Begin to train") for i_iter, batch in enumerate(trainloader): images, labels, _, _ = batch images = Variable(images).cuda() optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) pred = interp(model(images)) loss = loss_calc(pred, labels) loss.backward() optimizer.step() print('iter = ', i_iter, 'of', args.num_steps, 'completed, loss = ', loss.data.cpu().numpy()) logger.write("\n%d\t\t%.5f" % (i_iter, loss.data.cpu().numpy())) logger.flush() if i_iter >= args.num_steps - 1: print('save model ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC12_scenes_' + str(args.num_steps) + '.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, 'VOC12_scenes_' + str(i_iter) + '.pth')) logger.close() end = timeit.default_timer() print(end - start, 'seconds')
def main(): global opt, model, netContent opt = parser.parse_args() print(opt) gpuid = 0 cuda = True if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") opt.seed = random.randint(1, 10000) print("Random Seed: ", opt.seed) torch.manual_seed(opt.seed) if cuda: torch.cuda.manual_seed(opt.seed) cudnn.benchmark = True if opt.vgg_loss: print('===> Loading VGG model') netVGG = models.vgg19() netVGG.load_state_dict( model_zoo.load_url( 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth')) class _content_model(nn.Module): def __init__(self): super(_content_model, self).__init__() self.feature = nn.Sequential( *list(netVGG.features.children())[:-1]) def forward(self, x): out = self.feature(x) return out netContent = _content_model() print("===> Building model") #deeplab deeplab_res = Res_Deeplab(num_classes=21) saved_state_dict = torch.load('model/VOC12_scenes_20000.pth') deeplab_res.load_state_dict(saved_state_dict) deeplab_res = deeplab_res.eval() #deeplab_res.eval() mid = mid_layer() #SRResNet print("===> Building model") model = Net() finetune = True if finetune == True: model_pretrained = torch.load( 'model/model_DIV2K_noBN_96_epoch_36.pth', map_location=lambda storage, loc: storage)["model"] index = 0 for (src, dst) in zip(model_pretrained.parameters(), model.parameters()): if index != 0: list(model.parameters())[index].data = src.data index = index + 1 criterion = nn.MSELoss(size_average=False) print("===> Setting GPU") if cuda: model = model.cuda(gpuid) mid = mid.cuda(gpuid) deeplab_res = deeplab_res.cuda(gpuid) criterion = criterion.cuda(gpuid) if opt.vgg_loss: netContent = netContent.cuda(gpuid) # optionally resume from a checkpoint if opt.resume: if os.path.isfile(opt.resume): print("=> loading checkpoint '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) opt.start_epoch = checkpoint["epoch"] + 1 model.load_state_dict(checkpoint["model"].state_dict()) else: print("=> no checkpoint found at '{}'".format(opt.resume)) # optionally copy weights from a checkpoint if opt.pretrained: if os.path.isfile(opt.pretrained): print("=> loading model '{}'".format(opt.pretrained)) weights = torch.load(opt.pretrained) model.load_state_dict(weights['model'].state_dict()) else: print("=> no model found at '{}'".format(opt.pretrained)) print("===> Setting Optimizer") optimizer = optim.Adam(model.parameters(), lr=opt.lr) print("===> Training1") #root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5' root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train1_HDF5' files_num = len(os.listdir(root_dir)) for epoch in range(opt.start_epoch, opt.nEpochs + 1): #save_checkpoint(model, epoch) print("===> Loading datasets") x = random.sample(os.listdir(root_dir), files_num) for index in range(0, files_num): train_path = os.path.join(root_dir, x[index]) print("===> Training datasets: '{}'".format(train_path)) train_set = DatasetFromHdf5_OLD(train_path) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True) avgloss = train(training_data_loader, optimizer, deeplab_res, model, mid, criterion, epoch, gpuid) if epoch % 2 == 0: save_checkpoint(model, epoch)
def main(): """Create the model and start the training.""" h, w = map(int, args.input_size.split(',')) input_size = (h, w) print(input_size) cudnn.enabled = True gpu = args.gpu # Create network. model = Res_Deeplab(num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. ''' # 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 == 21 or not i_parts[1]=='layer5': # new_params['.'.join(i_parts[1:])] = saved_state_dict[i] # model.load_state_dict(new_params) ''' #model.float() #model.eval() # use_global_stats = True # model = nn.DataParallel(model) model.train() model.cuda() # model.cuda(args.gpu) cudnn.benchmark = True if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) # trainloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, max_iters=args.num_steps*args.batch_size, crop_size=input_size, # scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN), # batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True) dataset = GenericDataset(DATA_DIRECTORY, 'train', train_transform, mask_transform) trainloader = data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True) # optimizer = optim.SGD([{'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }, # {'params': get_10x_lr_params(model), 'lr': 10*args.learning_rate}], # lr=args.learning_rate, momentum=args.momentum,weight_decay=args.weight_decay) optimizer = optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), 1e-4) # optimizer = optim.Adam([{'params': get_1x_lr_params_NOscale(model), 'lr': args.learning_rate }], 1e-4) # optimizer = optim.Adam(get_1x_lr_params_NOscale(model), 1e-4) # optimizer = optim.Adam(get_10x_lr_params(model), lr=1e-5) optimizer.zero_grad() interp = nn.Upsample(size=input_size, mode='bilinear') for e in range(epochs): for i_iter, batch in enumerate(trainloader): images, labels, _, _ = batch # images = Variable(images).cuda(args.gpu) images = Variable(images).cuda() optimizer.zero_grad() adjust_learning_rate(optimizer, i_iter) pred = interp(model(images)) if i_iter % 50 == 0: vis.show( F.sigmoid(pred)[:, 1].cpu().data.round().numpy()[0:], labels.numpy()) loss = loss_calc(pred, labels.squeeze()) # loss = loss_calc(pred, labels.squeeze(), args.gpu) loss.backward() optimizer.step() print('iter = ', i_iter, 'of', args.num_steps, 'completed, loss = ', loss.data.cpu().numpy()) # if i_iter >= args.num_steps-1: # print('save model ...') # torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC12_scenes_'+str(args.num_steps)+'.pth')) # break if i_iter % 200 == 0 and i_iter != 0: print('taking snapshot ...') torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'VOC12_scenes_' + str(i_iter) + '.pth')) end = timeit.default_timer() print(end - start, 'seconds')