def main(): logger.auto_set_dir() global args parser = argparse.ArgumentParser() parser.add_argument('--dataroot', default='/home/hutao/lab/pytorchgo/example/ROAD/data', help='Path to source dataset') parser.add_argument('--batchSize', type=int, default=1, help='input batch size') parser.add_argument('--max_epoch', type=int, default=max_epoch, help='Number of training iterations') parser.add_argument('--optimizer', type=str, default='Adam', help='Optimizer to use | SGD, Adam') parser.add_argument('--lr', type=float, default=base_lr, help='learning rate') parser.add_argument('--momentum', type=float, default=0.99, help='Momentum for SGD') parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. default=0.5') parser.add_argument('--weight_decay', type=float, default=0.0005, help='Weight decay') parser.add_argument('--model', type=str, default='vgg16') parser.add_argument('--gpu', type=int, default=1) args = parser.parse_args() print(args) gpu = args.gpu os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) cuda = torch.cuda.is_available() torch.manual_seed(1337) if cuda: logger.info("random seed 1337") torch.cuda.manual_seed(1337) # Defining data loaders kwargs = { 'num_workers': 4, 'pin_memory': True, 'drop_last': True } if cuda else {} train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA( 'SYNTHIA', args.dataroot, split='train', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes( 'cityscapes', args.dataroot, split='val', transform=True, image_size=image_size), batch_size=1, shuffle=False) target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes( 'cityscapes', args.dataroot, split='train', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=True) if cuda: torch.set_default_tensor_type('torch.cuda.FloatTensor') if args.model == "vgg16": model = origin_model = torchfcn.models.Seg_model(n_class=class_num) vgg16 = torchfcn.models.VGG16(pretrained=True) model.copy_params_from_vgg16(vgg16) model_fix = torchfcn.models.Seg_model(n_class=class_num) model_fix.copy_params_from_vgg16(vgg16) for param in model_fix.parameters(): param.requires_grad = False elif args.model == "deeplabv2": # TODO may have problem! model = origin_model = torchfcn.models.Res_Deeplab( num_classes=class_num, image_size=image_size) saved_state_dict = model_zoo.load_url(Deeplabv2_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 class_num == 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_fix = torchfcn.models.Res_Deeplab(num_classes=class_num, image_size=image_size) model_fix.load_state_dict(new_params) else: raise ValueError("only support vgg16, deeplabv2!") netD = torchfcn.models.Domain_classifer(reverse=True) netD.apply(weights_init) model_summary([model, netD]) if cuda: model = model.cuda() netD = netD.cuda() # Defining optimizer if args.optimizer == 'SGD': raise ValueError("SGD is not prepared well..") optim = torch.optim.SGD([ { 'params': get_parameters(model, bias=False) }, { 'params': get_parameters(model, bias=True), 'lr': args.lr * 2, 'weight_decay': args.weight_decay }, ], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'Adam': if args.model == "vgg16": optim = torch.optim.Adam([ { 'params': get_parameters(model, bias=False), 'weight_decay': args.weight_decay }, { 'params': get_parameters(model, bias=True), 'lr': args.lr * 2, 'weight_decay': args.weight_decay }, ], lr=args.lr, betas=(args.beta1, 0.999)) elif args.model == "deeplabv2": optim = torch.optim.Adam(origin_model.optim_parameters(args.lr), lr=args.lr, betas=(args.beta1, 0.999), weight_decay=args.weight_decay) else: raise else: raise ValueError('Invalid optmizer argument. Has to be SGD or Adam') optimD = torch.optim.Adam(netD.parameters(), lr=dis_lr, weight_decay=args.weight_decay, betas=(0.7, 0.999)) optimizer_summary([optim, optimD]) trainer = MyTrainer_ROAD(cuda=cuda, model=model, model_fix=model_fix, netD=netD, optimizer=optim, optimizerD=optimD, train_loader=train_loader, target_loader=target_loader, val_loader=val_loader, batch_size=args.batchSize, image_size=image_size, loss_print_interval=LOSS_PRINT_INTERVAL) trainer.epoch = 0 trainer.iteration = 0 trainer.train()
def _init_(): if not os.path.exists('checkpoints'): os.makedirs('checkpoints') if not os.path.exists('checkpoints/' + args.exp_name): os.makedirs('checkpoints/' + args.exp_name) if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'): os.makedirs('checkpoints/' + args.exp_name + '/' + 'models') os.system('cp main.py checkpoints' + '/' + args.exp_name + '/' + 'main.py.backup') os.system('cp models/models.py checkpoints' + '/' + args.exp_name + '/' + 'models.py.backup') _init_() logger.auto_set_dir() logger.info(str(args)) args.cuda = not args.no_cuda and torch.cuda.is_available() torch.manual_seed(args.seed) if args.cuda: logger.info('Using GPU : ' + str(torch.cuda.current_device()) + ' from ' + str(torch.cuda.device_count()) + ' devices') torch.cuda.manual_seed(args.seed) else: logger.info('Using CPU') def train_batch(model, data): [enc_nn, metric_nn, softmax_module] = model
def train(args): logger.auto_set_dir() os.environ['CUDA_VISIBLE_DEVICES'] = '3' # Setup Augmentations data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()]) # Setup Dataloader data_loader = get_loader(args.dataset) data_path = get_data_path(args.dataset) t_loader = data_loader(data_path, is_transform=True, img_size=(args.img_rows, args.img_cols), epoch_scale=4, augmentations=data_aug, img_norm=args.img_norm) v_loader = data_loader(data_path, is_transform=True, split='val', img_size=(args.img_rows, args.img_cols), img_norm=args.img_norm) n_classes = t_loader.n_classes trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True) valloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=8) # Setup Metrics running_metrics = runningScore(n_classes) # Setup Model from model_zoo.deeplabv1 import VGG16_LargeFoV model = VGG16_LargeFoV(class_num=n_classes, image_size=[args.img_cols, args.img_rows], pretrained=True) #model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) model.cuda() # Check if model has custom optimizer / loss if hasattr(model, 'optimizer'): logger.warn("don't have customzed optimizer, use default setting!") optimizer = model.module.optimizer else: optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizer_summary(optimizer) if args.resume is not None: if os.path.isfile(args.resume): logger.info( "Loading model and optimizer from checkpoint '{}'".format( args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) logger.info("Loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: logger.info("No checkpoint found at '{}'".format(args.resume)) best_iou = -100.0 for epoch in tqdm(range(args.n_epoch), total=args.n_epoch): model.train() for i, (images, labels) in tqdm(enumerate(trainloader), total=len(trainloader), desc="training epoch {}/{}".format( epoch, args.n_epoch)): cur_iter = i + epoch * len(trainloader) cur_lr = adjust_learning_rate(optimizer, args.l_rate, cur_iter, args.n_epoch * len(trainloader), power=0.9) #if i > 10:break images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) #print(np.unique(outputs.data[0].cpu().numpy())) loss = CrossEntropyLoss2d_Seg(input=outputs, target=labels, class_num=n_classes) loss.backward() optimizer.step() if (i + 1) % 100 == 0: logger.info("Epoch [%d/%d] Loss: %.4f, lr: %.7f" % (epoch + 1, args.n_epoch, loss.data[0], cur_lr)) model.eval() for i_val, (images_val, labels_val) in tqdm(enumerate(valloader), total=len(valloader), desc="validation"): images_val = Variable(images_val.cuda(), volatile=True) labels_val = Variable(labels_val.cuda(), volatile=True) outputs = model(images_val) pred = outputs.data.max(1)[1].cpu().numpy() gt = labels_val.data.cpu().numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): logger.info("{}: {}".format(k, v)) running_metrics.reset() if score['Mean IoU : \t'] >= best_iou: best_iou = score['Mean IoU : \t'] state = { 'epoch': epoch + 1, 'mIoU': best_iou, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), } torch.save(state, os.path.join(logger.get_logger_dir(), "best_model.pkl"))
def main(): logger.auto_set_dir() parser = argparse.ArgumentParser() parser.add_argument('--dataroot', default='/home/hutao/lab/pytorchgo/example/ROAD/data', help='Path to source dataset') parser.add_argument('--batchSize', type=int, default=1, help='input batch size') parser.add_argument('--num_iters', type=int, default=100000, help='Number of training iterations') parser.add_argument('--optimizer', type=str, default='Adam', help='Optimizer to use | SGD, Adam') parser.add_argument('--lr', type=float, default=1.0e-5, help='learning rate') parser.add_argument('--momentum', type=float, default=0.99, help='Momentum for SGD') parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. default=0.5') parser.add_argument('--weight_decay', type=float, default=0.0005, help='Weight decay') parser.add_argument( '--interval_validate', type=int, default=3000, help= 'Period for validation. Model is validated every interval_validate iterations' ) parser.add_argument('--gpu', type=int, default=2) args = parser.parse_args() print(args) gpu = args.gpu os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) cuda = torch.cuda.is_available() torch.manual_seed(1337) if cuda: logger.info("random seed 1337") torch.cuda.manual_seed(1337) # Defining data loaders kwargs = { 'num_workers': 4, 'pin_memory': True, 'drop_last': True } if cuda else {} train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA( 'SYNTHIA', args.dataroot, split='train', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes( 'cityscapes', args.dataroot, split='val', transform=True, image_size=image_size), batch_size=1, shuffle=False) target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes( 'cityscapes', args.dataroot, split='train', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=True) # Defining models start_epoch = 0 start_iteration = 0 if cuda: torch.set_default_tensor_type('torch.cuda.FloatTensor') model = torchfcn.models.Seg_model(n_class=class_num) model_fix = torchfcn.models.Seg_model(n_class=class_num) for param in model_fix.parameters(): param.requires_grad = False netD = torchfcn.models.Domain_classifer(reverse=True) netD.apply(weights_init) model_summary(model_fix) model_summary(netD) vgg16 = torchfcn.models.VGG16(pretrained=True) model.copy_params_from_vgg16(vgg16) if cuda: model = model.cuda() netD = netD.cuda() # Defining optimizer if args.optimizer == 'SGD': optim = torch.optim.SGD([ { 'params': get_parameters(model, bias=False) }, { 'params': get_parameters(model, bias=True), 'lr': args.lr * 2, 'weight_decay': args.weight_decay }, ], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'Adam': optim = torch.optim.Adam([ { 'params': get_parameters(model, bias=False) }, { 'params': get_parameters(model, bias=True), 'lr': args.lr * 2 }, ], lr=args.lr, betas=(args.beta1, 0.999)) else: raise ValueError('Invalid optmizer argument. Has to be SGD or Adam') optimD = torch.optim.Adam(netD.parameters(), lr=0.0001, betas=(0.7, 0.999)) trainer = MyTrainer_ROAD(cuda=cuda, model=model, model_fix=model_fix, netD=netD, optimizer=optim, optimizerD=optimD, train_loader=train_loader, target_loader=target_loader, val_loader=val_loader, max_iter=args.num_iters, batch_size=args.batchSize, interval_validate=args.interval_validate, image_size=image_size) trainer.epoch = start_epoch trainer.iteration = start_iteration trainer.train()
def train(args): logger.auto_set_dir() from pytorchgo.utils.pytorch_utils import set_gpu set_gpu(args.gpu) # Setup Dataloader from pytorchgo.augmentation.segmentation import SubtractMeans, PIL2NP, RGB2BGR, PIL_Scale, Value255to0, ToLabel from torchvision.transforms import Compose, Normalize, ToTensor img_transform = Compose([ # notice the order!!! PIL_Scale(train_img_shape, Image.BILINEAR), PIL2NP(), RGB2BGR(), SubtractMeans(), ToTensor(), ]) label_transform = Compose([ PIL_Scale(train_img_shape, Image.NEAREST), PIL2NP(), Value255to0(), ToLabel() ]) val_img_transform = Compose([ PIL_Scale(train_img_shape, Image.BILINEAR), PIL2NP(), RGB2BGR(), SubtractMeans(), ToTensor(), ]) val_label_transform = Compose([ PIL_Scale(train_img_shape, Image.NEAREST), PIL2NP(), ToLabel(), # notice here, training, validation size difference, this is very tricky. ]) from pytorchgo.dataloader.pascal_voc_loader import pascalVOCLoader as common_voc_loader train_loader = common_voc_loader(split="train_aug", epoch_scale=1, img_transform=img_transform, label_transform=label_transform) validation_loader = common_voc_loader(split='val', img_transform=val_img_transform, label_transform=val_label_transform) n_classes = train_loader.n_classes trainloader = data.DataLoader(train_loader, batch_size=args.batch_size, num_workers=8, shuffle=True) valloader = data.DataLoader(validation_loader, batch_size=args.batch_size, num_workers=8) # Setup Metrics running_metrics = runningScore(n_classes) # Setup Model from pytorchgo.model.deeplabv1 import VGG16_LargeFoV from pytorchgo.model.deeplab_resnet import Res_Deeplab model = Res_Deeplab(NoLabels=n_classes, pretrained=True, output_all=False) from pytorchgo.utils.pytorch_utils import model_summary, optimizer_summary model_summary(model) def get_validation_miou(model): model.eval() for i_val, (images_val, labels_val) in tqdm(enumerate(valloader), total=len(valloader), desc="validation"): if i_val > 5 and is_debug == 1: break if i_val > 200 and is_debug == 2: break #img_large = torch.Tensor(np.zeros((1, 3, 513, 513))) #img_large[:, :, :images_val.shape[2], :images_val.shape[3]] = images_val output = model(Variable(images_val, volatile=True).cuda()) output = output pred = output.data.max(1)[1].cpu().numpy() #pred = output[:, :images_val.shape[2], :images_val.shape[3]] gt = labels_val.numpy() running_metrics.update(gt, pred) score, class_iou = running_metrics.get_scores() for k, v in score.items(): logger.info("{}: {}".format(k, v)) running_metrics.reset() return score['Mean IoU : \t'] model.cuda() # Check if model has custom optimizer / loss if hasattr(model, 'optimizer'): logger.warn("don't have customzed optimizer, use default setting!") optimizer = model.module.optimizer else: optimizer = torch.optim.SGD(model.optimizer_params(args.l_rate), lr=args.l_rate, momentum=0.99, weight_decay=5e-4) optimizer_summary(optimizer) if args.resume is not None: if os.path.isfile(args.resume): logger.info( "Loading model and optimizer from checkpoint '{}'".format( args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) logger.info("Loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: logger.info("No checkpoint found at '{}'".format(args.resume)) best_iou = 0 logger.info('start!!') for epoch in tqdm(range(args.n_epoch), total=args.n_epoch): model.train() for i, (images, labels) in tqdm(enumerate(trainloader), total=len(trainloader), desc="training epoch {}/{}".format( epoch, args.n_epoch)): if i > 10 and is_debug == 1: break if i > 200 and is_debug == 2: break cur_iter = i + epoch * len(trainloader) cur_lr = adjust_learning_rate(optimizer, args.l_rate, cur_iter, args.n_epoch * len(trainloader), power=0.9) images = Variable(images.cuda()) labels = Variable(labels.cuda()) optimizer.zero_grad() outputs = model(images) # use fusion score loss = CrossEntropyLoss2d_Seg(input=outputs, target=labels, class_num=n_classes) #for i in range(len(outputs) - 1): #for i in range(1): # loss = loss + CrossEntropyLoss2d_Seg(input=outputs[i], target=labels, class_num=n_classes) loss.backward() optimizer.step() if (i + 1) % 100 == 0: logger.info( "Epoch [%d/%d] Loss: %.4f, lr: %.7f, best mIoU: %.7f" % (epoch + 1, args.n_epoch, loss.data[0], cur_lr, best_iou)) cur_miou = get_validation_miou(model) if cur_miou >= best_iou: best_iou = cur_miou state = { 'epoch': epoch + 1, 'mIoU': best_iou, 'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict(), } torch.save(state, os.path.join(logger.get_logger_dir(), "best_model.pth"))
args.mixup_alpha = 0.4 else: args.mixup_alpha = 1.5 else: if not args.manimixup: args.mixup_alpha = 1.0 else: args.mixup_alpha = 2.0 if (args.pointmixup and args.rot) and (not args.manimixup): args.align = True else: args.align = False # logger.auto_set_dir('d', "{}_{}".format(args.data, args.model)) logger.auto_set_dir('d', "{}".format(args.savename)) if args.data == 'modelnet40': num_class = 40 dataset_cls = ModelNet40Cls else: raise NotImplementedError args.epochs = int(args.epochs) from tensorboardX import SummaryWriter writer = SummaryWriter(comment=args.savename) writer.add_text('args', str(args), 0) transforms_test = d_utils.PointcloudToTensor() if args.rot:
def main(): logger.auto_set_dir() global args, best_prec1 import argparse parser = argparse.ArgumentParser(description="PyTorch implementation of Temporal Segment Networks") parser.add_argument('--dataset', type=str,default="something", choices=['something', 'jester', 'moments']) parser.add_argument('--modality', type=str, default="RGB", choices=['RGB', 'Flow']) parser.add_argument('--train_list', type=str, default="") parser.add_argument('--val_list', type=str, default="") parser.add_argument('--root_path', type=str, default="") parser.add_argument('--store_name', type=str, default="") # ========================= Model Configs ========================== parser.add_argument('--arch', type=str, default="BNInception") parser.add_argument('--num_segments', type=int, default=3) parser.add_argument('--consensus_type', type=str, default='avg') parser.add_argument('--k', type=int, default=3) parser.add_argument('--dropout', '--do', default=0.8, type=float, metavar='DO', help='dropout ratio (default: 0.5)') parser.add_argument('--loss_type', type=str, default="nll", choices=['nll']) parser.add_argument('--img_feature_dim', default=256, type=int, help="the feature dimension for each frame") # ========================= Learning Configs ========================== parser.add_argument('--epochs', default=120, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('-b', '--batch_size', default=128, type=int, metavar='N', help='mini-batch size (default: 256)') parser.add_argument('--lr', '--learning-rate', default=0.001, type=float, metavar='LR', help='initial learning rate') parser.add_argument('--lr_steps', default=[50, 100], type=float, nargs="+", metavar='LRSteps', help='epochs to decay learning rate by 10') parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') parser.add_argument('--weight-decay', '--wd', default=5e-4, type=float, metavar='W', help='weight decay (default: 5e-4)') parser.add_argument('--clip-gradient', '--gd', default=20, type=float, metavar='W', help='gradient norm clipping (default: disabled)') parser.add_argument('--no_partialbn', '--npb', default=False, action="store_true") # ========================= Monitor Configs ========================== parser.add_argument('--print-freq', '-p', default=20, type=int, metavar='N', help='print frequency (default: 10)') parser.add_argument('--eval-freq', '-ef', default=5, type=int, metavar='N', help='evaluation frequency (default: 5)') # ========================= Runtime Configs ========================== parser.add_argument('-j', '--workers', default=30, type=int, metavar='N', help='number of data loading workers (default: 4)') parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set') parser.add_argument('--snapshot_pref', type=str, default="") parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)') parser.add_argument('--gpu', type=str, default='4') parser.add_argument('--flow_prefix', default="", type=str) parser.add_argument('--root_log', type=str, default='log') parser.add_argument('--root_model', type=str, default='model') parser.add_argument('--root_output', type=str, default='output') args = parser.parse_args() args.consensus_type = "TRN" os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu device_ids = [int(id) for id in args.gpu.split(',')] assert len(device_ids) >1, "TRN must run with GPU_num > 1" args.root_log = logger.get_logger_dir() args.root_model = logger.get_logger_dir() args.root_output = logger.get_logger_dir() categories, args.train_list, args.val_list, args.root_path, prefix = datasets_video.return_dataset(args.dataset, args.modality) num_class = len(categories) args.store_name = '_'.join(['TRN', args.dataset, args.modality, args.arch, args.consensus_type, 'segment%d'% args.num_segments]) print('storing name: ' + args.store_name) model = TSN(num_class, args.num_segments, args.modality, base_model=args.arch, consensus_type=args.consensus_type, dropout=args.dropout, img_feature_dim=args.img_feature_dim, partial_bn=not args.no_partialbn) crop_size = model.crop_size scale_size = model.scale_size input_mean = model.input_mean input_std = model.input_std policies = model.get_optim_policies() train_augmentation = model.get_augmentation() if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model)#TODO, , device_ids=[int(id) for id in args.gpu.split(',')] if torch.cuda.is_available(): model.cuda() if args.resume: if os.path.isfile(args.resume): print(("=> loading checkpoint '{}'".format(args.resume))) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) print(("=> loaded checkpoint '{}' (epoch {})" .format(args.evaluate, checkpoint['epoch']))) else: print(("=> no checkpoint found at '{}'".format(args.resume))) cudnn.benchmark = True # Data loading code if args.modality != 'RGBDiff': normalize = GroupNormalize(input_mean, input_std) else: normalize = IdentityTransform() if args.modality == 'RGB': data_length = 1 elif args.modality in ['Flow', 'RGBDiff']: data_length = 5 train_loader = torch.utils.data.DataLoader( TSNDataSet(args.root_path, args.train_list, num_segments=args.num_segments, new_length=data_length, modality=args.modality, image_tmpl=prefix, transform=torchvision.transforms.Compose([ train_augmentation, Stack(roll=(args.arch in ['BNInception','InceptionV3'])), ToTorchFormatTensor(div=(args.arch not in ['BNInception','InceptionV3'])), normalize, ])), batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader( TSNDataSet(args.root_path, args.val_list, num_segments=args.num_segments, new_length=data_length, modality=args.modality, image_tmpl=prefix, random_shift=False, transform=torchvision.transforms.Compose([ GroupScale(int(scale_size)), GroupCenterCrop(crop_size), Stack(roll=(args.arch in ['BNInception','InceptionV3'])), ToTorchFormatTensor(div=(args.arch not in ['BNInception','InceptionV3'])), normalize, ])), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) # define loss function (criterion) and optimizer if args.loss_type == 'nll': criterion = torch.nn.CrossEntropyLoss().cuda() else: raise ValueError("Unknown loss type") for group in policies: logger.info('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format( group['name'], len(group['params']), group['lr_mult'], group['decay_mult'])) optimizer = torch.optim.SGD(policies, args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if args.evaluate: validate(val_loader, model, criterion, 0) return log_training = open(os.path.join(args.root_log, '%s.csv' % args.store_name), 'w') for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch, args.lr_steps) # train for one epoch train(train_loader, model, criterion, optimizer, epoch, log_training) # evaluate on validation set if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1: prec1 = validate(val_loader, model, criterion, (epoch + 1) * len(train_loader), log_training) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, }, is_best)
def main(): logger.auto_set_dir() parser = argparse.ArgumentParser() parser.add_argument( '--dataroot', default='/home/hutao/lab/pytorchgo/example/LSD-seg/data', help='Path to source dataset') parser.add_argument('--batchSize', type=int, default=1, help='input batch size') parser.add_argument('--num_iters', type=int, default=100000, help='Number of training iterations') parser.add_argument('--optimizer', type=str, default='Adam', help='Optimizer to use | SGD, Adam') parser.add_argument('--lr', type=float, default=1.0e-5, help='learning rate') parser.add_argument('--momentum', type=float, default=0.99, help='Momentum for SGD') parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. default=0.5') parser.add_argument('--weight_decay', type=float, default=0.0005, help='Weight decay') parser.add_argument( '--interval_validate', type=int, default=500, help= 'Period for validation. Model is validated every interval_validate iterations' ) parser.add_argument( '--resume', default='', help= "path to the current checkpoint for resuming training. Do not specify if model has to be trained from scratch" ) parser.add_argument('--method', default='LSD', help="Method to use for training | LSD, sourceonly") parser.add_argument('--l1_weight', type=float, default=1, help='L1 weight') parser.add_argument('--adv_weight', type=float, default=0.1, help='Adv_weight') parser.add_argument('--c_weight', type=float, default=0.1, help='C_weight') parser.add_argument('--gpu', type=int, default=0) args = parser.parse_args() print(args) gpu = args.gpu out = logger.get_logger_dir() resume = args.resume os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) cuda = torch.cuda.is_available() torch.manual_seed(1337) if cuda: torch.cuda.manual_seed(1337) # Defining data loaders image_size = [640, 320] kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {} train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA( 'SYNTHIA', args.dataroot, split='train', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA( 'SYNTHIA', args.dataroot, split='val', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=False, **kwargs) target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes( 'cityscapes', args.dataroot, split='train', transform=True, image_size=image_size), batch_size=args.batchSize, shuffle=True) # Defining models start_epoch = 0 start_iteration = 0 if args.method == 'sourceonly': model = torchfcn.models.FCN8s_sourceonly(n_class=class_num) elif args.method == 'LSD': model = torchfcn.models.FCN8s_LSD(n_class=class_num) netG = torchfcn.models._netG() netD = torchfcn.models._netD() netD.apply(weights_init) netG.apply(weights_init) else: raise ValueError('method argument can be either sourceonly or LSD') if resume: checkpoint = torch.load(resume) model.load_state_dict(checkpoint['model_state_dict']) start_epoch = checkpoint['epoch'] start_iteration = checkpoint['iteration'] else: vgg16 = torchfcn.models.VGG16(pretrained=True) model.copy_params_from_vgg16(vgg16) if cuda: model = model.cuda() if args.method == 'LSD': netD = netD.cuda() netG = netG.cuda() # Defining optimizer if args.optimizer == 'SGD': optim = torch.optim.SGD([ { 'params': get_parameters(model, bias=False) }, { 'params': get_parameters(model, bias=True), 'lr': args.lr * 2, 'weight_decay': args.weight_decay }, ], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'Adam': optim = torch.optim.Adam([ { 'params': get_parameters(model, bias=False) }, { 'params': get_parameters(model, bias=True), 'lr': args.lr * 2 }, ], lr=args.lr, betas=(args.beta1, 0.999)) else: raise ValueError('Invalid optmizer argument. Has to be SGD or Adam') if args.method == 'LSD': optimD = torch.optim.Adam(netD.parameters(), lr=0.0001, betas=(0.7, 0.999)) optimG = torch.optim.Adam(netG.parameters(), lr=0.0001, betas=(0.7, 0.999)) if resume: optim.load_state_dict(checkpoint['optim_state_dict']) # Defining trainer object, and start training if args.method == 'sourceonly': trainer = torchfcn.Trainer_sourceonly( cuda=cuda, model=model, optimizer=optim, train_loader=train_loader, target_loader=target_loader, val_loader=val_loader, out=out, max_iter=args.num_iters, interval_validate=args.interval_validate) trainer.epoch = start_epoch trainer.iteration = start_iteration trainer.train() elif args.method == 'LSD': trainer = Trainer_LSD(cuda=cuda, model=model, netD=netD, netG=netG, optimizer=optim, optimizerD=optimD, optimizerG=optimG, train_loader=train_loader, target_loader=target_loader, l1_weight=args.l1_weight, adv_weight=args.adv_weight, c_weight=args.c_weight, val_loader=val_loader, out=out, max_iter=args.num_iters, interval_validate=args.interval_validate, image_size=image_size) trainer.epoch = start_epoch trainer.iteration = start_iteration trainer.train()