def main(): start_epoch = 0 args = parse_args() # Use GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu if args.gpu != '' else str( opt.gpu_id) use_gpu = torch.cuda.is_available() and (args.gpu != '' or int(opt.gpu_id)) >= 0 if not os.path.isdir(opt.checkpoint): os.makedirs(opt.checkpoint) # Data print('==> Preparing dataset') input_dim = opt.input_size train_transformer = TrainTransform(size=input_dim) test_transformer = TestTransform(size=input_dim) try: if isinstance(opt.trainset, list): datalist = [] for dataset, freq, max_skip in zip(opt.trainset, opt.datafreq, opt.max_skip): ds = DATA_CONTAINER[dataset]( train=True, sampled_frames=opt.sampled_frames, transform=train_transformer, max_skip=max_skip, samples_per_video=opt.samples_per_video) datalist += [ds] * freq trainset = data.ConcatDataset(datalist) else: max_skip = opt.max_skip[0] if isinstance(opt.max_skip, list) else opt.max_skip trainset = DATA_CONTAINER[opt.trainset]( train=True, sampled_frames=opt.sampled_frames, transform=train_transformer, max_skip=max_skip, samples_per_video=opt.samples_per_video) except KeyError as ke: print( '[ERROR] invalide dataset name is encountered. The current acceptable datasets are:' ) print(list(DATA_CONTAINER.keys())) exit() testset = DATA_CONTAINER[opt.valset](train=False, transform=test_transformer, samples_per_video=1) trainloader = data.DataLoader(trainset, batch_size=opt.train_batch, shuffle=True, num_workers=opt.workers, collate_fn=multibatch_collate_fn) testloader = data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=opt.workers, collate_fn=multibatch_collate_fn) # Model print("==> creating model") net = STAN(opt.keydim, opt.valdim) print(' Total params: %.2fM' % (sum(p.numel() for p in net.parameters()) / 1000000.0)) net.eval() if use_gpu: net = net.cuda() # set training parameters for p in net.parameters(): p.requires_grad = True criterion = None celoss = cross_entropy_loss if opt.loss == 'ce': criterion = celoss elif opt.loss == 'iou': criterion = mask_iou_loss elif opt.loss == 'both': criterion = lambda pred, target, obj: celoss( pred, target, obj) + mask_iou_loss(pred, target, obj) else: raise TypeError('unknown training loss %s' % opt.loss) optimizer = None if opt.solver == 'sgd': optimizer = optim.SGD(net.parameters(), lr=opt.learning_rate, momentum=opt.momentum[0], weight_decay=opt.weight_decay) elif opt.solver == 'adam': optimizer = optim.Adam(net.parameters(), lr=opt.learning_rate, betas=opt.momentum, weight_decay=opt.weight_decay) else: raise TypeError('unkown solver type %s' % opt.solver) # Resume title = 'STAN' minloss = float('inf') opt.checkpoint = osp.join(osp.join(opt.checkpoint, opt.valset)) if not osp.exists(opt.checkpoint): os.mkdir(opt.checkpoint) if opt.resume: # Load checkpoint. print('==> Resuming from checkpoint {}'.format(opt.resume)) assert os.path.isfile( opt.resume), 'Error: no checkpoint directory found!' # opt.checkpoint = os.path.dirname(opt.resume) checkpoint = torch.load(opt.resume) minloss = checkpoint['minloss'] start_epoch = checkpoint['epoch'] net.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) skips = checkpoint['max_skip'] try: if isinstance(skips, list): for idx, skip in enumerate(skips): trainloader.dataset.datasets[idx].set_max_skip(skip) else: trainloader.dataset.set_max_skip(skip) except: print('[Warning] Initializing max skip fail') logger = Logger(os.path.join(opt.checkpoint, opt.mode + '_log.txt'), resume=True) else: if opt.initial: print('==> Initialize model with weight file {}'.format( opt.initial)) weight = torch.load(opt.initial) if isinstance(weight, OrderedDict): net.load_param(weight) else: net.load_param(weight['state_dict']) logger = Logger(os.path.join(opt.checkpoint, opt.mode + '_log.txt'), resume=False) start_epoch = 0 logger.set_items(['Epoch', 'LR', 'Train Loss']) # Train and val for epoch in range(start_epoch): adjust_learning_rate(optimizer, epoch, opt) for epoch in range(start_epoch, opt.epochs): print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, opt.epochs, opt.learning_rate)) adjust_learning_rate(optimizer, epoch, opt) train_loss = train(trainloader, model=net, criterion=criterion, optimizer=optimizer, epoch=epoch, use_cuda=use_gpu, iter_size=opt.iter_size, mode=opt.mode, threshold=opt.iou_threshold) if (epoch + 1) % opt.epoch_per_test == 0: test_loss = test(testloader, model=net, criterion=criterion, epoch=epoch, use_cuda=use_gpu, opt=opt) # append logger file logger.log(epoch + 1, opt.learning_rate, train_loss) # adjust max skip if (epoch + 1) % opt.epochs_per_increment == 0: if isinstance(trainloader.dataset, data.ConcatDataset): for dataset in trainloader.dataset.datasets: dataset.increase_max_skip() else: trainloader.dataset.increase_max_skip() # save model is_best = train_loss <= minloss minloss = min(minloss, train_loss) skips = [ds.max_skip for ds in trainloader.dataset.datasets] \ if isinstance(trainloader.dataset, data.ConcatDataset) \ else trainloader.dataset.max_skip save_checkpoint( { 'epoch': epoch + 1, 'state_dict': net.state_dict(), 'loss': train_loss, 'minloss': minloss, 'optimizer': optimizer.state_dict(), 'max_skip': skips, }, epoch + 1, is_best, checkpoint=opt.checkpoint, filename=opt.mode) logger.close() print('minimum loss:') print(minloss)
def main(): settings_print_interval = 1 # How often to print loss and other info settings_batch_size = 4 # Batch size 80 default 64 settings_num_workers = 16 # Number of workers for image loading settings_normalize_mean = [0.485, 0.456, 0.406] # Normalize mean (default pytorch ImageNet values) settings_normalize_std = [0.229, 0.224, 0.225] # Normalize std (default pytorch ImageNet values) settings_search_area_factor = 4.0 # Image patch size relative to target size settings_feature_sz = 24 # Size of feature map settings_output_sz = settings_feature_sz * 16 # Size of input image patches 24*16 settings_segm_use_distance = True # Settings for the image sample and proposal generation settings_center_jitter_factor = {'train': 0, 'test1': 1.5, 'test2': 1.5} settings_scale_jitter_factor = {'train': 0, 'test1': 0.25, 'test2': 0.25} #################################################################################################### start_epoch = 0 random.seed(0) args = parse_args() # Use GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu if args.gpu != '' else str(opt.gpu_id) use_gpu = torch.cuda.is_available() and (args.gpu != '' or int(opt.gpu_id)) >= 0 gpu_ids = [int(val) for val in args.gpu.split(',')] if not os.path.isdir(opt.checkpoint): os.makedirs(opt.checkpoint) # Data print('==> Preparing dataset') input_size = opt.input_size train_transformer = TrainTransform(size=input_size) #train_transformer = TrainTransform_Noresize() test_transformer = TestTransform(size=input_size) try: if isinstance(opt.trainset, list): datalist = [] for dataset, freq, max_skip in zip(opt.trainset, opt.datafreq, opt.max_skip): ds = DATA_CONTAINER[dataset]( train=True, sampled_frames=opt.sampled_frames, transform=train_transformer, max_skip=max_skip, samples_per_video=opt.samples_per_video ) datalist += [ds] * freq trainset = data.ConcatDataset(datalist) else: max_skip = opt.max_skip[0] if isinstance(opt.max_skip, list) else opt.max_skip trainset = DATA_CONTAINER[opt.trainset]( train=True, sampled_frames=opt.sampled_frames, transform=train_transformer, max_skip=max_skip, samples_per_video=opt.samples_per_video ) except KeyError as ke: print('[ERROR] invalide dataset name is encountered. The current acceptable datasets are:') print(list(DATA_CONTAINER.keys())) exit() testset = DATA_CONTAINER[opt.valset]( train=False, transform=test_transformer, samples_per_video=1 ) trainloader = data.DataLoader(trainset, batch_size=opt.train_batch, shuffle=True, num_workers=opt.workers, collate_fn=multibatch_collate_fn, drop_last=True) testloader = data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=opt.workers, collate_fn=multibatch_collate_fn) ######################################################################################### vos_train = Vos(split='train') transform_train = torchvision.transforms.Compose([dltransforms.ToTensorAndJitter(0.2), torchvision.transforms.Normalize(mean=settings_normalize_mean, std=settings_normalize_std)]) data_processing_train = segm_processing.SegmProcessing(search_area_factor=settings_search_area_factor, output_sz=settings_output_sz, center_jitter_factor=settings_center_jitter_factor, scale_jitter_factor=settings_scale_jitter_factor, mode='pair', transform=transform_train, use_distance=settings_segm_use_distance) dataset_train = segm_sampler.SegmSampler([vos_train], [1], samples_per_epoch=1000 * settings_batch_size * 8, max_gap=50, processing=data_processing_train) loader_train = LTRLoader('train', dataset_train, training=True, batch_size=settings_batch_size, num_workers=settings_num_workers, shuffle=True, drop_last=True, stack_dim=1) ######################################################################################### # Model print("==> creating model") net = AMB(opt.keydim, opt.valdim, 'train', mode=opt.mode, iou_threshold=opt.iou_threshold) print(' Total params: %.2fM' % (sum(p.numel() for p in net.parameters())/1000000.0)) net.eval() if use_gpu: net = net.cuda() assert opt.train_batch % len(gpu_ids) == 0 net = nn.DataParallel(net, device_ids=gpu_ids, dim=0) # set training parameters #for p in net.parameters(): # p.requires_grad = True for name, param in net.named_parameters(): #print(name) if 'Encoder' in name: param.requires_grad = False # 冻结 backbone 梯度 else: param.requires_grad = True criterion = None celoss = cross_entropy_loss if opt.loss == 'ce': criterion = celoss elif opt.loss == 'iou': criterion = mask_iou_loss elif opt.loss == 'both': criterion = lambda pred, target, obj: celoss(pred, target, obj) + mask_iou_loss(pred, target, obj) else: raise TypeError('unknown training loss %s' % opt.loss) optimizer = None if opt.solver == 'sgd': optimizer = optim.SGD(net.parameters(), lr=opt.learning_rate, momentum=opt.momentum[0], weight_decay=opt.weight_decay) elif opt.solver == 'adam': optimizer = optim.Adam(net.parameters(), lr=opt.learning_rate, betas=opt.momentum, weight_decay=opt.weight_decay) else: raise TypeError('unkown solver type %s' % opt.solver) # Resume title = 'Appearance Memory Bank' minloss = float('inf') opt.checkpoint = osp.join(osp.join(opt.checkpoint, opt.valset)) if not osp.exists(opt.checkpoint): os.mkdir(opt.checkpoint) if opt.resume: # Load checkpoint. print('==> Resuming from checkpoint {}'.format(opt.resume)) assert os.path.isfile(opt.resume), 'Error: no checkpoint directory found!' # opt.checkpoint = os.path.dirname(opt.resume) checkpoint = torch.load(opt.resume) minloss = checkpoint['minloss'] start_epoch = checkpoint['epoch'] net.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) skips = checkpoint['max_skip'] try: if isinstance(skips, list): for idx, skip in enumerate(skips): trainloader.dataset.datasets[idx].set_max_skip(skip) else: trainloader.dataset.set_max_skip(skip) except: print('[Warning] Initializing max skip fail') logger = Logger(os.path.join(opt.checkpoint, opt.mode+'_log.txt'), resume=True) else: if opt.initial: print('==> Initialize model with weight file {}'.format(opt.initial)) weight = torch.load(opt.initial) if isinstance(weight, OrderedDict): net.module.load_param(weight) else: net.module.load_param(weight['state_dict']) logger = Logger(os.path.join(opt.checkpoint, opt.mode+'_log.txt'), resume=False) start_epoch = 0 logger.set_items(['Epoch', 'LR', 'Train Loss']) # Train and val for epoch in range(start_epoch): adjust_learning_rate(optimizer, epoch, opt) for epoch in range(start_epoch, opt.epochs): print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, opt.epochs, opt.learning_rate)) adjust_learning_rate(optimizer, epoch, opt) net.module.phase = 'train' train_loss = train(loader_train, # loader_train trainloader model=net, criterion=criterion, optimizer=optimizer, epoch=epoch, use_cuda=use_gpu, iter_size=opt.iter_size, mode=opt.mode, threshold=opt.iou_threshold) if (epoch + 1) % opt.epoch_per_test == 0: net.module.phase = 'test' test_loss = test(testloader, model=net.module, criterion=criterion, epoch=epoch, use_cuda=use_gpu) # append logger file logger.log(epoch+1, opt.learning_rate, train_loss) # adjust max skip if (epoch + 1) % opt.epochs_per_increment == 0: if isinstance(trainloader.dataset, data.ConcatDataset): for dataset in trainloader.dataset.datasets: dataset.increase_max_skip() else: trainloader.dataset.increase_max_skip() # save model is_best = train_loss <= minloss minloss = min(minloss, train_loss) skips = [ds.max_skip for ds in trainloader.dataset.datasets] \ if isinstance(trainloader.dataset, data.ConcatDataset) \ else trainloader.dataset.max_skip save_checkpoint({ 'epoch': epoch + 1, 'state_dict': net.state_dict(), 'loss': train_loss, 'minloss': minloss, 'optimizer': optimizer.state_dict(), 'max_skip': skips, }, epoch + 1, is_best, checkpoint=opt.checkpoint, filename=opt.mode) logger.close() print('minimum loss:') print(minloss)