def main(): """Create the model and start the evaluation process.""" args = get_arguments() os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu gpus = [int(i) for i in args.gpu.split(',')] h, w = map(int, args.input_size.split(',')) input_size = (h, w) model = Res_Deeplab(num_classes=args.num_classes) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) lip_dataset = LIPDataSet(args.data_dir, 'test', crop_size=input_size, transform=transform) num_samples = len(lip_dataset) valloader = data.DataLoader(lip_dataset, batch_size=args.batch_size * len(gpus), shuffle=False, pin_memory=True) restore_from = args.restore_from state_dict = model.state_dict().copy() state_dict_old = torch.load(restore_from) for key, nkey in zip(state_dict_old.keys(), state_dict.keys()): if key != nkey: # remove the 'module.' in the 'key' state_dict[key[7:]] = deepcopy(state_dict_old[key]) else: state_dict[key] = deepcopy(state_dict_old[key]) model.load_state_dict(state_dict) model.eval() model.cuda() parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus)) write_results(parsing_preds, scales, centers, args.data_dir, 'test', args.save_dir, input_size=input_size)
def main(): """Create the model and start the training.""" if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) writer = SummaryWriter(args.snapshot_dir) gpus = [int(i) for i in args.gpu.split(',')] if not args.gpu == 'None': os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu h, w = map(int, args.input_size.split(',')) input_size = [h, w] cudnn.enabled = True # cudnn related setting cudnn.benchmark = True torch.backends.cudnn.deterministic = False torch.backends.cudnn.enabled = True deeplab = Res_Deeplab(num_classes=args.num_classes) # dump_input = torch.rand((args.batch_size, 3, input_size[0], input_size[1])) # writer.add_graph(deeplab.cuda(), dump_input.cuda(), verbose=False) saved_state_dict = torch.load(args.restore_from) new_params = deeplab.state_dict().copy() for i in saved_state_dict: i_parts = i.split('.') # print(i_parts) if not i_parts[0] == 'fc': new_params['.'.join(i_parts[0:])] = saved_state_dict[i] deeplab.load_state_dict(new_params) model = DataParallelModel(deeplab) model.cuda() criterion = CriterionAll() criterion = DataParallelCriterion(criterion) criterion.cuda() normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) trainloader = data.DataLoader(LIPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform), batch_size=args.batch_size * len(gpus), shuffle=True, num_workers=2, pin_memory=True) #lip_dataset = LIPDataSet(args.data_dir, 'val', crop_size=input_size, transform=transform) #num_samples = len(lip_dataset) #valloader = data.DataLoader(lip_dataset, batch_size=args.batch_size * len(gpus), # shuffle=False, pin_memory=True) optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() total_iters = args.epochs * len(trainloader) for epoch in range(args.start_epoch, args.epochs): model.train() for i_iter, batch in enumerate(trainloader): i_iter += len(trainloader) * epoch lr = adjust_learning_rate(optimizer, i_iter, total_iters) images, labels, edges, _ = batch labels = labels.long().cuda(non_blocking=True) edges = edges.long().cuda(non_blocking=True) preds = model(images) loss = criterion(preds, [labels, edges]) optimizer.zero_grad() loss.backward() optimizer.step() if i_iter % 100 == 0: writer.add_scalar('learning_rate', lr, i_iter) writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter) if i_iter % 500 == 0: images_inv = inv_preprocess(images, args.save_num_images) labels_colors = decode_parsing(labels, args.save_num_images, args.num_classes, is_pred=False) edges_colors = decode_parsing(edges, args.save_num_images, 2, is_pred=False) if isinstance(preds, list): preds = preds[0] preds_colors = decode_parsing(preds[0][-1], args.save_num_images, args.num_classes, is_pred=True) pred_edges = decode_parsing(preds[1][-1], args.save_num_images, 2, is_pred=True) img = vutils.make_grid(images_inv, normalize=False, scale_each=True) lab = vutils.make_grid(labels_colors, normalize=False, scale_each=True) pred = vutils.make_grid(preds_colors, normalize=False, scale_each=True) edge = vutils.make_grid(edges_colors, normalize=False, scale_each=True) pred_edge = vutils.make_grid(pred_edges, normalize=False, scale_each=True) writer.add_image('Images/', img, i_iter) writer.add_image('Labels/', lab, i_iter) writer.add_image('Preds/', pred, i_iter) writer.add_image('Edges/', edge, i_iter) writer.add_image('PredEdges/', pred_edge, i_iter) print('iter = {} of {} completed, loss = {}'.format( i_iter, total_iters, loss.data.cpu().numpy())) torch.save( model.state_dict(), osp.join(args.snapshot_dir, 'LIP_epoch_' + str(epoch) + '.pth')) #parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus)) #mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size) #print(mIoU) #writer.add_scalars('mIoU', mIoU, epoch) end = timeit.default_timer() print(end - start, 'seconds')
def main(): """Create the model and start the evaluation process.""" args = get_arguments() update_config(config, args) print(args) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu gpus = [int(i) for i in args.gpu.split(',')] h, w = map(int, args.input_size.split(',')) input_size = (h, w) model = get_cls_net(config=config, num_classes=args.num_classes, is_train=False) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) print('-------Load Data', args.data_dir) if 'vehicle_parsing_dataset' in args.data_dir: parsing_dataset = VPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform, list_path=args.list_path) elif 'LIP' in args.data_dir: parsing_dataset = LIPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform) num_samples = len(parsing_dataset) valloader = data.DataLoader(parsing_dataset, batch_size=args.batch_size * len(gpus), shuffle=False, pin_memory=True) print('-------Load Weight', args.restore_from) restore_from = args.restore_from state_dict = model.state_dict().copy() state_dict_old = torch.load(restore_from) state_dict_old = state_dict_old['state_dict'] for key, nkey in zip(state_dict_old.keys(), state_dict.keys()): if key != nkey: # remove the 'module.' in the 'key' state_dict[key[7:]] = deepcopy(state_dict_old[key]) else: state_dict[key] = deepcopy(state_dict_old[key]) model.load_state_dict(state_dict) model = DataParallelModel(model) model.eval() model.cuda() print('-------Start Testing...') parsing_preds, scales, centers, time_list = valid(model, valloader, input_size, num_samples, len(gpus)) print('-------Saving Results', args.save_dir) write_results(parsing_preds, scales, centers, args.data_dir, args.dataset, args.save_dir, input_size=input_size) print('total time is ', sum(time_list)) print('avg time is ', sum(time_list) / len(time_list))
def main(): """Create the model and start the training.""" print(args) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) writer = SummaryWriter(args.snapshot_dir) gpus = [int(i) for i in args.gpu.split(',')] if not args.gpu == 'None': os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu h, w = map(int, args.input_size.split(',')) input_size = [h, w] # cudnn related setting cudnn.enabled = True cudnn.benchmark = True torch.backends.cudnn.deterministic = False torch.backends.cudnn.enabled = True deeplab = get_cls_net(config=config, is_train=True) # model = DataParallelModel(deeplab) print('-------Load Weight', args.restore_from) saved_state_dict = torch.load(args.restore_from) if args.start_epoch > 0: model = DataParallelModel(deeplab) model.load_state_dict(saved_state_dict['state_dict']) else: new_params = deeplab.state_dict().copy() state_dict_pretrain = saved_state_dict for state_name in state_dict_pretrain: if state_name in new_params: new_params[state_name] = state_dict_pretrain[state_name] print('LOAD', state_name) else: print('NOT LOAD', state_name) deeplab.load_state_dict(new_params) model = DataParallelModel(deeplab) print('-------Load Weight', args.restore_from) model.cuda() criterion = CriterionAll(loss_type=args.loss) criterion = DataParallelCriterion(criterion) criterion.cuda() normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) print("-------Loading data...") if 'vehicle_parsing_dataset' in args.data_dir: parsing_dataset = VPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform) elif 'LIP' in args.data_dir: parsing_dataset = LIPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform) print("Data dir : ", args.data_dir) print("Dataset : ", args.dataset) trainloader = data.DataLoader(parsing_dataset, batch_size=args.batch_size * len(gpus), shuffle=True, num_workers=8, pin_memory=True) ''' list_map = [] for part in deeplab.path_list: list_map = list_map + list(map(id, part.parameters())) base_params = filter(lambda p: id(p) not in list_map, deeplab.parameters()) params_list = [] params_list.append({'params': base_params, 'lr':args.learning_rate*0.1}) for part in deeplab.path_list: params_list.append({'params': part.parameters()}) print ('len(params_list)',len(params_list)) ''' optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) if args.start_epoch > 0: optimizer.load_state_dict(saved_state_dict['optimizer']) print('========Load Optimizer', args.restore_from) total_iters = args.epochs * len(trainloader) for epoch in range(args.start_epoch, args.epochs): model.train() for i_iter, batch in enumerate(trainloader): i_iter += len(trainloader) * epoch lr = adjust_learning_rate(optimizer, i_iter, total_iters) images, labels, _ = batch labels = labels.long().cuda(non_blocking=True) preds = model(images) loss = criterion(preds, labels) optimizer.zero_grad() loss.backward() optimizer.step() if i_iter % 100 == 0: writer.add_scalar('learning_rate', lr, i_iter) writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter) print( f'epoch = {epoch}, iter = {i_iter}/{total_iters}, lr={lr:.6f}, loss = {loss.data.cpu().numpy():.6f}' ) if (epoch + 1) % args.save_step == 0 or epoch == args.epochs: time.sleep(10) print("-------Saving checkpoint...") save_checkpoint(model, epoch, optimizer) time.sleep(10) save_checkpoint(model, epoch, optimizer) 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) h, w = map(int, args.input_size.split(',')) input_size = [h, w] best_f1 = 0 torch.cuda.set_device(args.local_rank) try: world_size = int(os.environ['WORLD_SIZE']) distributed = world_size > 1 except: distributed = False world_size = 1 if distributed: dist.init_process_group(backend=args.dist_backend, init_method='env://') rank = 0 if not distributed else dist.get_rank() writer = SummaryWriter(osp.join(args.snapshot_dir, TIMESTAMP)) if rank == 0 else None normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) if args.type == 'Helen': train_dataset = HelenDataSet('dataset/Helen_align_with_hair', args.dataset, crop_size=input_size, transform=transform) val_dataset = HelenDataSet('dataset/Helen_align_with_hair', 'test', crop_size=input_size, transform=transform) args.num_classes = 11 elif args.type == 'LaPa': train_dataset = LapaDataset('dataset/LaPa/origin', args.dataset, crop_size=input_size, transform=transform) val_dataset = LapaDataset('dataset/LaPa/origin', 'test', crop_size=input_size, transform=transform) args.num_classes = 11 elif args.type == 'Celeb': train_dataset = CelebAMaskHQDataSet('dataset/CelebAMask-HQ', args.dataset, crop_size=input_size, transform=transform) val_dataset = CelebAMaskHQDataSet('dataset/CelebAMask-HQ', 'test', crop_size=input_size, transform=transform) args.num_classes = 19 elif args.type == 'LIP': train_dataset = LIPDataSet('dataset/LIP', args.dataset, crop_size=input_size, transform=transform) val_dataset = LIPDataSet('dataset/LIP', 'val', crop_size=input_size, transform=transform) args.num_classes = 20 if distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=2, pin_memory=True, drop_last=True, sampler=train_sampler) num_samples = len(val_dataset) valloader = data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=True, drop_last=False) cudnn.enabled = True # cudnn related setting cudnn.benchmark = True torch.backends.cudnn.deterministic = False torch.backends.cudnn.enabled = True if distributed: model = AGRNet(args.num_classes) else: model = AGRNet(args.num_classes, InPlaceABN) if args.restore_from is not None: model.load_state_dict( torch.load(args.restore_from, map_location='cuda:{}'.format(args.local_rank)), True) else: resnet_params = torch.load( os.path.join(args.snapshot_dir, 'resnet101-imagenet.pth')) new_params = model.state_dict().copy() for i in resnet_params: i_parts = i.split('.') # print(i_parts) if not i_parts[0] == 'fc': new_params['.'.join(i_parts[0:])] = resnet_params[i] model.load_state_dict(new_params) model.cuda() if distributed: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) else: model = SingleGPU(model) # CriterionCrossEntropyEdgeParsing_boundary_agrnet_loss for AGRNet, CriterionCrossEntropyEdgeParsing_boundary_eagrnet_loss for EAGRNet criterion = CriterionCrossEntropyEdgeParsing_boundary_agrnet_loss( loss_weight=[args.l1, args.l2, args.l3, args.l4], num_classes=args.num_classes) criterion.cuda() optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() total_iters = args.epochs * len(trainloader) for epoch in range(args.start_epoch, args.epochs): model.train() if distributed: train_sampler.set_epoch(epoch) for i_iter, batch in enumerate(trainloader): i_iter += len(trainloader) * epoch lr = adjust_learning_rate(optimizer, i_iter, total_iters) images, labels, edges, _ = batch labels = labels.long().cuda(non_blocking=True) edges = edges.long().cuda(non_blocking=True) preds = model(images) loss = criterion(preds, [labels, edges]) optimizer.zero_grad() loss.backward() optimizer.step() with torch.no_grad(): loss = loss.detach() * labels.shape[0] count = labels.new_tensor([labels.shape[0]], dtype=torch.long) if dist.is_initialized(): dist.all_reduce(count, dist.ReduceOp.SUM) dist.all_reduce(loss, dist.ReduceOp.SUM) loss /= count.item() if not dist.is_initialized() or dist.get_rank() == 0: if i_iter % 50 == 0: writer.add_scalar('learning_rate', lr, i_iter) writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter) if i_iter % 500 == 0: images_inv = inv_preprocess(images, args.save_num_images) labels_colors = decode_parsing(labels, args.save_num_images, args.num_classes, is_pred=False) edges_colors = decode_parsing(edges, args.save_num_images, 2, is_pred=False) if isinstance(preds, list): preds = preds[0] preds_colors = decode_parsing(preds[0], args.save_num_images, args.num_classes, is_pred=True) pred_edges = decode_parsing(preds[1], args.save_num_images, 2, is_pred=True) img = vutils.make_grid(images_inv, normalize=False, scale_each=True) lab = vutils.make_grid(labels_colors, normalize=False, scale_each=True) pred = vutils.make_grid(preds_colors, normalize=False, scale_each=True) edge = vutils.make_grid(edges_colors, normalize=False, scale_each=True) pred_edge = vutils.make_grid(pred_edges, normalize=False, scale_each=True) writer.add_image('Images/', img, i_iter) writer.add_image('Labels/', lab, i_iter) writer.add_image('Preds/', pred, i_iter) writer.add_image('Edge/', edge, i_iter) writer.add_image('Pred_edge/', pred_edge, i_iter) print('iter = {} of {} completed, loss = {}'.format( i_iter, total_iters, loss.data.cpu().numpy())) if not dist.is_initialized() or dist.get_rank() == 0: save_path = os.path.join(args.data_dir, TIMESTAMP) if not os.path.exists(save_path): os.makedirs(save_path) parsing_preds, scales, centers = valid( model, valloader, input_size, num_samples, osp.join(args.snapshot_dir, save_path)) mIoU, f1 = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, val_dataset, input_size, 'test', True, type=args.type) if f1['mean'] > best_f1: torch.save(model.module.state_dict(), osp.join(args.snapshot_dir, TIMESTAMP, 'best.pth')) best_f1 = f1['mean'] print(mIoU) print(f1) writer.add_scalars('mIoU', mIoU, epoch) writer.add_scalars('f1', f1, epoch) if epoch % args.test_fre == 0: torch.save( model.module.state_dict(), osp.join(args.snapshot_dir, TIMESTAMP, 'epoch_' + str(epoch) + '.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) timestramp = args.date writer = SummaryWriter(os.path.join(args.snapshot_dir, timestramp)) gpus = [int(i) for i in args.gpu.split(',')] if not args.gpu == 'None': os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu h, w = map(int, args.input_size.split(',')) input_size = [h, w] cudnn.enabled = True # cudnn related setting cudnn.benchmark = True torch.backends.cudnn.deterministic = False torch.backends.cudnn.enabled = True deeplab = Res_Deeplab(num_classes=args.num_classes) # dump_input = torch.rand((args.batch_size, 3, input_size[0], input_size[1])) # writer.add_graph(deeplab.cuda(), dump_input.cuda(), verbose=False) model = DataParallelModel(deeplab) if args.resume: # when restore form the same network, it is useful here checkpoint = torch.load(args.restore_from) model.load_state_dict(checkpoint['net']) args.start_epoch = checkpoint['epoch'] else: saved_state_dict = torch.load(args.restore_from) new_params = deeplab.state_dict().copy() for i in saved_state_dict: i_parts = i.split('.') if not i_parts[0] == 'fc': new_params['.'.join(i_parts[0:])] = saved_state_dict[i] deeplab.load_state_dict(new_params) model.cuda() criterion = CriterionAll() criterion = DataParallelCriterion(criterion) criterion.cuda() normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) trainloader = data.DataLoader(LIPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform), batch_size=args.batch_size * len(gpus), shuffle=True, num_workers=2, pin_memory=True) lip_dataset = LIPDataSet(args.data_dir, 'val', crop_size=input_size, transform=transform) num_samples = len(lip_dataset) valloader = data.DataLoader(lip_dataset, batch_size=args.batch_size * len(gpus), shuffle=False, pin_memory=True) optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) optimizer.zero_grad() total_iters = args.epochs * len(trainloader) log = Logger(os.path.join(args.log_dir, '{}_train.log'.format(timestramp)), level='debug') for epoch in range(args.start_epoch, args.epochs): model.train() #lr = adjust_learning_rate_pose(optimizer, epoch) for i_iter, batch in enumerate(trainloader): i_iter += len(trainloader) * epoch lr = adjust_learning_rate(optimizer, i_iter, total_iters) images, labels, r1, r2, r3, r4, l0, l1, l2, l3, l4, l5, _ = batch labels = labels.long().cuda(non_blocking=True) r1 = r1.long().cuda(non_blocking=True) r2 = r2.long().cuda(non_blocking=True) r3 = r3.long().cuda(non_blocking=True) r4 = r4.long().cuda(non_blocking=True) l0 = l0.long().cuda(non_blocking=True) l1 = l1.long().cuda(non_blocking=True) l2 = l2.long().cuda(non_blocking=True) l3 = l3.long().cuda(non_blocking=True) l4 = l4.long().cuda(non_blocking=True) l5 = l5.long().cuda(non_blocking=True) preds = model(images) loss = criterion( preds, [[labels], [r1, r2, r3, r4], [l0, l1, l2, l3, l4, l5]]) optimizer.zero_grad() loss.backward() optimizer.step() if i_iter % 100 == 0: writer.add_scalar('learning_rate', lr, i_iter) writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter) if i_iter % 10 == 0: log.logger.info( 'epoch = {} iter = {} of {} completed, lr = {}, loss = {}'. format(epoch, i_iter, total_iters, lr, loss.data.cpu().numpy())) parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus)) mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size) log.logger.info('epoch = {}'.format(epoch)) log.logger.info(str(mIoU)) writer.add_scalars('mIoU', mIoU, epoch) # save the model snapshot state = {"net": model.module.state_dict(), "epoch": epoch} torch.save( state, osp.join(args.snapshot_dir, timestramp, 'LIP_epoch_' + str(epoch) + '.pth')) end = timeit.default_timer() print(end - start, 'seconds')
def main(): """Create the model and start the training.""" print(args) if not os.path.exists(args.snapshot_dir): os.makedirs(args.snapshot_dir) writer = SummaryWriter(args.snapshot_dir) gpus = [int(i) for i in args.gpu.split(',')] if not args.gpu == 'None': os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu h, w = map(int, args.input_size.split(',')) input_size = [h, w] cudnn.enabled = True # cudnn related setting cudnn.benchmark = True torch.backends.cudnn.deterministic = False torch.backends.cudnn.enabled = True deeplab = Res_Deeplab(num_classes=args.num_classes) # dump_input = torch.rand((args.batch_size, 3, input_size[0], input_size[1])) # writer.add_graph(deeplab.cuda(), dump_input.cuda(), verbose=False) saved_state_dict = torch.load(args.restore_from) if args.start_epoch > 0: model = DataParallelModel(deeplab) model.load_state_dict(saved_state_dict['state_dict']) else: new_params = deeplab.state_dict().copy() state_dict_pretrain = saved_state_dict['state_dict'] for i in state_dict_pretrain: splits = i.split('.') state_name = '.'.join(splits[1:]) if state_name in new_params: new_params[state_name] = state_dict_pretrain[i] else: print('NOT LOAD', state_name) deeplab.load_state_dict(new_params) model = DataParallelModel(deeplab) print('-------Load Weight', args.restore_from) model.cuda() criterion = CriterionAll() criterion = DataParallelCriterion(criterion) criterion.cuda() normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) trainloader = data.DataLoader(LIPDataSet(args.data_dir, args.dataset, crop_size=input_size, transform=transform), batch_size=args.batch_size * len(gpus), shuffle=True, num_workers=4, pin_memory=True) num_samples = 5000 ''' list_map = [] for part in deeplab.path_list: list_map = list_map + list(map(id, part.parameters())) base_params = filter(lambda p: id(p) not in list_map, deeplab.parameters()) params_list = [] params_list.append({'params': base_params, 'lr':args.learning_rate*0.1}) for part in deeplab.path_list: params_list.append({'params': part.parameters()}) print ('len(params_list)',len(params_list)) ''' optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) if args.start_epoch >= 0: optimizer.load_state_dict(saved_state_dict['optimizer']) print('========Load Optimizer', args.restore_from) total_iters = args.epochs * len(trainloader) for epoch in range(args.start_epoch, args.epochs): model.train() for i_iter, batch in enumerate(trainloader): i_iter += len(trainloader) * epoch lr = adjust_learning_rate(optimizer, i_iter, total_iters) images, labels, _ = batch labels = labels.long().cuda(non_blocking=True) preds = model(images) loss = criterion(preds, labels) optimizer.zero_grad() loss.backward() optimizer.step() if i_iter % 100 == 0: writer.add_scalar('learning_rate', lr, i_iter) writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter) # if i_iter % 500 == 0: # images_inv = inv_preprocess(images, args.save_num_images) # labels_colors = decode_parsing(labels, args.save_num_images, args.num_classes, is_pred=False) # edges_colors = decode_parsing(edges, args.save_num_images, 2, is_pred=False) # if isinstance(preds, list): # preds = preds[0] # preds_colors = decode_parsing(preds[0][-1], args.save_num_images, args.num_classes, is_pred=True) # pred_edges = decode_parsing(preds[1][-1], args.save_num_images, 2, is_pred=True) # img = vutils.make_grid(images_inv, normalize=False, scale_each=True) # lab = vutils.make_grid(labels_colors, normalize=False, scale_each=True) # pred = vutils.make_grid(preds_colors, normalize=False, scale_each=True) # edge = vutils.make_grid(edges_colors, normalize=False, scale_each=True) # pred_edge = vutils.make_grid(pred_edges, normalize=False, scale_each=True) # writer.add_image('Images/', img, i_iter) # writer.add_image('Labels/', lab, i_iter) # writer.add_image('Preds/', pred, i_iter) # writer.add_image('Edges/', edge, i_iter) # writer.add_image('PredEdges/', pred_edge, i_iter) print('epoch = {}, iter = {} of {} completed,lr={}, loss = {}'. format(epoch, i_iter, total_iters, lr, loss.data.cpu().numpy())) if epoch % 2 == 0 or epoch == args.epochs: time.sleep(10) save_checkpoint(model, epoch, optimizer) # parsing_preds, scales, centers = valid(model, valloader, input_size, num_samples, len(gpus)) # mIoU = compute_mean_ioU(parsing_preds, scales, centers, args.num_classes, args.data_dir, input_size) # print(mIoU) # writer.add_scalars('mIoU', mIoU, epoch) time.sleep(10) save_checkpoint(model, epoch, optimizer) end = timeit.default_timer() print(end - start, 'seconds')