def train(**kwargs): opt._parse(kwargs) opt.model_name = 'bfe_test' # set random seed and cudnn benchmark torch.manual_seed(opt.seed) os.makedirs(opt.save_dir, exist_ok=True) use_gpu = torch.cuda.is_available() sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt')) print('=========user config==========') pprint(opt._state_dict()) print('============end===============') if use_gpu: print('currently using GPU') cudnn.benchmark = True torch.cuda.manual_seed_all(opt.seed) else: print('currently using cpu') print('initializing dataset {}'.format(opt.dataset)) dataset = data_manager.init_dataset(name=opt.dataset, mode=opt.mode) pin_memory = True if use_gpu else False summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log')) trainloader = DataLoader(ImageData(dataset.train, TrainTransform(opt.datatype)), sampler=RandomIdentitySampler( dataset.train, opt.num_instances), batch_size=opt.train_batch, num_workers=opt.workers, pin_memory=pin_memory, drop_last=True) queryloader = DataLoader(ImageData(dataset.query, TestTransform(opt.datatype)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) galleryloader = DataLoader(ImageData(dataset.gallery, TestTransform(opt.datatype)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) queryFliploader = DataLoader(ImageData(dataset.query, TestTransform(opt.datatype, True)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) galleryFliploader = DataLoader(ImageData(dataset.gallery, TestTransform(opt.datatype, True)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) print('initializing model ...') model = BFE(dataset.num_train_pids, 1.0, 0.33) optim_policy = model.get_optim_policy() if opt.pretrained_model: state_dict = torch.load(opt.pretrained_model)['state_dict'] # state_dict = {k: v for k, v in state_dict.items() \ # if not ('reduction' in k or 'softmax' in k)} model.load_state_dict(state_dict, False) print('load pretrained model ' + opt.pretrained_model) print('model size: {:.5f}M'.format( sum(p.numel() for p in model.parameters()) / 1e6)) if use_gpu: model = nn.DataParallel(model).cuda() reid_evaluator = ResNetEvaluator(model) if opt.evaluate: reid_evaluator.evaluate(queryloader, galleryloader, queryFliploader, galleryFliploader, re_ranking=opt.re_ranking, savefig=opt.savefig) return # xent_criterion = nn.CrossEntropyLoss() xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids) if opt.loss == 'triplet': embedding_criterion = TripletLoss(opt.margin) elif opt.loss == 'lifted': embedding_criterion = LiftedStructureLoss(hard_mining=True) elif opt.loss == 'weight': embedding_criterion = Margin() def criterion(triplet_y, softmax_y, labels): losses = [embedding_criterion(output, labels)[0] for output in triplet_y] + \ [xent_criterion(output, labels) for output in softmax_y] loss = sum(losses) return loss # get optimizer if opt.optim == "sgd": optimizer = torch.optim.SGD(optim_policy, lr=opt.lr, momentum=0.9, weight_decay=opt.weight_decay) else: optimizer = torch.optim.Adam(optim_policy, lr=opt.lr, weight_decay=opt.weight_decay) start_epoch = opt.start_epoch # get trainer and evaluator reid_trainer = cls_tripletTrainer(opt, model, optimizer, criterion, summary_writer) def adjust_lr(optimizer, ep): if ep < 10: lr = opt.lr * 0.1 * (ep / 10.0) # warm_up elif ep < 50: lr = opt.lr * (ep // 5 + 1) elif ep < 200: lr = opt.lr * 10.0 elif ep < 300: lr = opt.lr else: lr = opt.lr * 0.1 for p in optimizer.param_groups: p['lr'] = lr # start training best_rank1 = opt.best_rank best_epoch = 0 for epoch in range(start_epoch, opt.max_epoch): if opt.adjust_lr: adjust_lr(optimizer, epoch + 1) reid_trainer.train(epoch, trainloader) # skip if not save model if opt.eval_step > 0 and (epoch + 1) % opt.eval_step == 0 or ( epoch + 1) == opt.max_epoch: if opt.mode == 'class': rank1 = test(model, queryloader) else: rank1 = reid_evaluator.evaluate(queryloader, galleryloader, queryFliploader, galleryFliploader) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch + 1 }, is_best=is_best, save_dir=opt.save_dir, filename='checkpoint_ep' + str(epoch + 1) + '.pth.tar') print('Best rank-1 {:.1%}, achived at epoch {}'.format( best_rank1, best_epoch))
def train(**kwargs): opt._parse(kwargs) # set random seed and cudnn benchmark torch.manual_seed(opt.seed) use_gpu = torch.cuda.is_available() sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt')) print('=========user config==========') pprint(opt._state_dict()) print('============end===============') if use_gpu: print('currently using GPU') cudnn.benchmark = True torch.cuda.manual_seed_all(opt.seed) else: print('currently using cpu') print('initializing dataset {}'.format(opt.dataset)) dataset = data_manager.init_dataset(name=opt.dataset) pin_memory = True if use_gpu else False summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log')) if 'triplet' in opt.model_name: trainloader = DataLoader( ImageData(dataset.train, TrainTransform(opt.height, opt.width)), sampler=RandomIdentitySampler(dataset.train, opt.num_instances), batch_size=opt.train_batch, num_workers=opt.workers, pin_memory=pin_memory, drop_last=True) else: trainloader = DataLoader(ImageData( dataset.train, TrainTransform(opt.height, opt.width)), batch_size=opt.train_batch, shuffle=True, num_workers=opt.workers, pin_memory=pin_memory) queryloader = DataLoader(ImageData(dataset.query, TestTransform(opt.height, opt.width)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) galleryloader = DataLoader(ImageData(dataset.gallery, TestTransform(opt.height, opt.width)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) print('initializing model ...') if opt.model_name == 'softmax' or opt.model_name == 'softmax_triplet': model, optim_policy = get_baseline_model(dataset.num_train_pids) elif opt.model_name == 'triplet': model, optim_policy = get_baseline_model(num_classes=None) print('model size: {:.5f}M'.format( sum(p.numel() for p in model.parameters()) / 1e6)) # xent_criterion = nn.CrossEntropyLoss() xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids) tri_criterion = TripletLoss(opt.margin) def cls_criterion(cls_scores, targets): cls_loss = xent_criterion(cls_scores, targets) return cls_loss def triplet_criterion(feat, targets): triplet_loss, _, _ = tri_criterion(feat, targets) return triplet_loss def cls_tri_criterion(cls_scores, feat, targets): cls_loss = xent_criterion(cls_scores, targets) triplet_loss, _, _ = tri_criterion(feat, targets) loss = cls_loss + triplet_loss return loss # get optimizer optimizer = torch.optim.Adam(optim_policy, lr=opt.lr, weight_decay=opt.weight_decay) def adjust_lr(optimizer, ep): if ep < 20: lr = 1e-4 * (ep + 1) / 2 elif ep < 80: lr = 1e-3 * opt.num_gpu elif ep < 180: lr = 1e-4 * opt.num_gpu elif ep < 300: lr = 1e-5 * opt.num_gpu elif ep < 320: lr = 1e-5 * 0.1**((ep - 320) / 80) * opt.num_gpu elif ep < 400: lr = 1e-6 elif ep < 480: lr = 1e-4 * opt.num_gpu else: lr = 1e-5 * opt.num_gpu for p in optimizer.param_groups: p['lr'] = lr start_epoch = opt.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() # get trainer and evaluator if opt.model_name == 'softmax': reid_trainer = clsTrainer(opt, model, optimizer, cls_criterion, summary_writer) elif opt.model_name == 'softmax_triplet': reid_trainer = cls_tripletTrainer(opt, model, optimizer, cls_tri_criterion, summary_writer) elif opt.model_name == 'triplet': reid_trainer = tripletTrainer(opt, model, optimizer, triplet_criterion, summary_writer) reid_evaluator = ResNetEvaluator(model) # start training best_rank1 = -np.inf best_epoch = 0 for epoch in range(start_epoch, opt.max_epoch): if opt.step_size > 0: adjust_lr(optimizer, epoch + 1) reid_trainer.train(epoch, trainloader) # skip if not save model if opt.eval_step > 0 and (epoch + 1) % opt.eval_step == 0 or ( epoch + 1) == opt.max_epoch: rank1 = reid_evaluator.evaluate(queryloader, galleryloader) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch + 1, }, is_best=is_best, save_dir=opt.save_dir, filename='checkpoint_ep' + str(epoch + 1) + '.pth.tar') print('Best rank-1 {:.1%}, achived at epoch {}'.format( best_rank1, best_epoch))
def main(): logger = logging.getLogger('global') global criterion_xent, criterion_triplet, criterion_center if os.path.exists(cfg.TRAIN.LOG_DIR): shutil.rmtree(cfg.TRAIN.LOG_DIR) os.makedirs(cfg.TRAIN.LOG_DIR) init_log('global', logging.INFO) # log add_file_handler('global', os.path.join(cfg.TRAIN.LOG_DIR, 'logs.txt'), logging.INFO) summary_writer = SummaryWriter(cfg.TRAIN.LOG_DIR) # visualise dataset, train_loader, _, _ = build_data_loader() model = BagReID_RESNET(dataset.num_train_bags) criterion_xent = CrossEntropyLabelSmooth(dataset.num_train_bags, use_gpu=cfg.CUDA) criterion_triplet = TripletLoss(margin=cfg.TRAIN.MARGIN) criterion_center = CenterLoss(dataset.num_train_bags, cfg.MODEL.GLOBAL_FEATS + cfg.MODEL.PART_FEATS, use_gpu=cfg.CUDA) if cfg.TRAIN.OPTIM == "sgd": optimizer = torch.optim.SGD(model.parameters(), lr=cfg.SOLVER.LEARNING_RATE, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY) else: optimizer = torch.optim.Adam(model.parameters(), lr=cfg.SOLVER.LEARNING_RATE, weight_decay=cfg.SOLVER.WEIGHT_DECAY) center_optimizer = torch.optim.SGD(criterion_center.parameters(), lr=cfg.SOLVER.LEARNING_RATE_CENTER) optimizers = [optimizer, center_optimizer] schedulers = build_lr_schedulers(optimizers) if cfg.CUDA: model.cuda() if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model, device_ids=cfg.DEVICES) logger.info("model prepare done") # start training for epoch in range(cfg.TRAIN.NUM_EPOCHS): train(epoch, train_loader, model, criterion, optimizers, summary_writer) for scheduler in schedulers: scheduler.step() # skip if not save model if cfg.TRAIN.EVAL_STEP > 0 and (epoch + 1) % cfg.TRAIN.EVAL_STEP == 0 \ or (epoch + 1) == cfg.TRAIN.NUM_EPOCHS: if cfg.CUDA and torch.cuda.device_count() > 1: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch + 1 }, is_best=False, save_dir=cfg.TRAIN.SNAPSHOT_DIR, filename='checkpoint_ep' + str(epoch + 1) + '.pth')
def train(**kwargs): opt._parse(kwargs) #opt.lr=0.00002 opt.model_name = 'AlignedReid' # set random seed and cudnn benchmark torch.manual_seed(opt.seed) os.makedirs(opt.save_dir, exist_ok=True) use_gpu = torch.cuda.is_available() sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt')) print('=========user config==========') pprint(opt._state_dict()) print('============end===============') if use_gpu: print('currently using GPU') cudnn.benchmark = True torch.cuda.manual_seed_all(opt.seed) else: print('currently using cpu') print('initializing dataset {}'.format(opt.dataset)) dataset = data_manager.init_dataset(name=opt.dataset, mode=opt.mode) pin_memory = True if use_gpu else False summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log')) trainloader = DataLoader(ImageData(dataset.train, TrainTransform(opt.datatype)), sampler=RandomIdentitySampler( dataset.train, opt.num_instances), batch_size=opt.train_batch, num_workers=opt.workers, pin_memory=pin_memory, drop_last=True) queryloader = DataLoader(ImageData(dataset.query, TestTransform(opt.datatype)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) galleryloader = DataLoader(ImageData(dataset.gallery, TestTransform(opt.datatype)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) queryFliploader = queryloader galleryFliploader = galleryloader # queryFliploader = DataLoader( # ImageData(dataset.query, TestTransform(opt.datatype, True)), # batch_size=opt.test_batch, num_workers=opt.workers, # pin_memory=pin_memory # ) # # galleryFliploader = DataLoader( # ImageData(dataset.gallery, TestTransform(opt.datatype, True)), # batch_size=opt.test_batch, num_workers=opt.workers, # pin_memory=pin_memory # ) print('initializing model ...') model = AlignedResNet50(num_classes=dataset.num_train_pids, loss={'softmax', 'metric'}, aligned=True, use_gpu=use_gpu) optim_policy = model.get_optim_policy() if opt.pretrained_model: state_dict = torch.load(opt.pretrained_model)['state_dict'] # state_dict = {k: v for k, v in state_dict.items() \ # if not ('reduction' in k or 'softmax' in k)} model.load_state_dict(state_dict, False) print('load pretrained model ' + opt.pretrained_model) print('model size: {:.5f}M'.format( sum(p.numel() for p in model.parameters()) / 1e6)) if use_gpu: model = nn.DataParallel(model).cuda() reid_evaluator = AlignedEvaluator(model) if opt.evaluate: #rank1 = test(model, queryloader, galleryloader, use_gpu) reid_evaluator.evaluate(queryloader, galleryloader, queryFliploader, galleryFliploader, re_ranking=opt.re_ranking, savefig=opt.savefig, test_distance='global') return # xent_criterion = nn.CrossEntropyLoss() xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids, use_gpu=use_gpu) embedding_criterion = TripletLossAlignedReID(margin=opt.margin) # def criterion(triplet_y, softmax_y, labels): # losses = [embedding_criterion(output, labels)[0] for output in triplet_y] + \ # [xent_criterion(output, labels) for output in softmax_y] # loss = sum(losses) # return loss def criterion(outputs, features, local_features, labels): if opt.htri_only: if isinstance(features, tuple): global_loss, local_loss = DeepSupervision( embedding_criterion, features, labels) else: global_loss, local_loss = embedding_criterion( features, labels, local_features) else: if isinstance(outputs, tuple): xent_loss = DeepSupervision(xent_criterion, outputs, labels) else: xent_loss = xent_criterion(outputs, labels) if isinstance(features, tuple): global_loss, local_loss = DeepSupervision( embedding_criterion, features, labels) else: global_loss, local_loss = embedding_criterion( features, labels, local_features) loss = xent_loss + global_loss + local_loss return loss # get optimizer if opt.optim == "sgd": optimizer = torch.optim.SGD(optim_policy, lr=opt.lr, momentum=0.9, weight_decay=opt.weight_decay) else: optimizer = torch.optim.Adam(optim_policy, lr=opt.lr, weight_decay=opt.weight_decay) start_epoch = opt.start_epoch # get trainer and evaluator reid_trainer = AlignedTrainer(opt, model, optimizer, criterion, summary_writer) def adjust_lr(optimizer, ep): if ep < 50: lr = opt.lr * (ep // 5 + 1) elif ep < 200: lr = opt.lr * 10 elif ep < 300: lr = opt.lr else: lr = opt.lr * 0.1 for p in optimizer.param_groups: p['lr'] = lr # start training best_rank1 = opt.best_rank best_epoch = 0 print('start train......') for epoch in range(start_epoch, opt.max_epoch): if opt.adjust_lr: adjust_lr(optimizer, epoch + 1) reid_trainer.train(epoch, trainloader) # skip if not save model if opt.eval_step > 0 and (epoch + 1) % opt.eval_step == 0 or ( epoch + 1) == opt.max_epoch: # just avoid out of memory during eval,and can't save the model if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch + 1 }, is_best=0, save_dir=opt.save_dir, filename='checkpoint_ep' + str(epoch + 1) + '.pth.tar') if opt.mode == 'class': rank1 = test(model, queryloader) else: rank1 = reid_evaluator.evaluate(queryloader, galleryloader, queryFliploader, galleryFliploader) #rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() if is_best: save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch + 1 }, is_best=is_best, save_dir=opt.save_dir, filename='checkpoint_ep' + str(epoch + 1) + '.pth.tar') print('Best rank-1 {:.1%}, achived at epoch {}'.format( best_rank1, best_epoch))
model.load_state_dict(state_dict, False) print('load pretrained model ' + opt.pretrained_model) print('model size: {:.5f}M'.format(sum(p.numel() for p in model.parameters()) / 1e6)) if use_gpu: model = nn.DataParallel(model).cuda() reid_evaluator = ResNetEvaluator(model) if opt.evaluate: reid_evaluator.evaluate(queryloader, galleryloader, queryFliploader, galleryFliploader, re_ranking=opt.re_ranking, savefig=opt.savefig) return #xent_criterion = nn.CrossEntropyLoss() #print(dataset.num_train_pids) xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids) xent_criterion2 = CrossEntropyLabelSmooth2(dataset.num_train_pids) if opt.loss == 'triplet': embedding_criterion = TripletLoss(opt.margin) elif opt.loss == 'lifted': embedding_criterion = LiftedStructureLoss(hard_mining=True) elif opt.loss == 'weight': embedding_criterion = Margin() def criterion(a ,b, c, d, e, f, triplet_y, softmax_y, labels): losses = [embedding_criterion(output, labels)[0] for output in triplet_y] + \ [xent_criterion(output, labels) for output in softmax_y]+\ [xent_criterion2(i, j, labels) for i, j in zip(a,a)]+\ [xent_criterion2(i, j, labels) for i, j in zip(b,b)]+\ [xent_criterion2(i, j, labels) for i, j in zip(c,c)]+\
def train(**kwargs): opt._parse(kwargs) ##设置程序的所有参数 # set random seed and cudnn benchmark torch.manual_seed(opt.seed) os.makedirs(opt.save_dir, exist_ok=True) use_gpu = torch.cuda.is_available() sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt')) print('=========user config==========') pprint(opt._state_dict()) print('============end===============') if use_gpu: print('currently using GPU') cudnn.benchmark = True torch.cuda.manual_seed_all(opt.seed) else: print('currently using cpu') print('initializing dataset {}'.format(opt.dataset)) dataset = data_manager.init_dataset(name=opt.dataset, mode=opt.mode) pin_memory = True if use_gpu else False trainloader = DataLoader(ImageData(dataset.train, TrainTransform(opt.datatype)), batch_size=opt.train_batch, num_workers=opt.workers, pin_memory=pin_memory, drop_last=True) queryloader = DataLoader(ImageData(dataset.query, TestTransform(opt.datatype)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) galleryloader = DataLoader(ImageData(dataset.gallery, TestTransform(opt.datatype)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) queryFliploader = DataLoader(ImageData(dataset.query, TestTransform(opt.datatype, True)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) galleryFliploader = DataLoader(ImageData(dataset.gallery, TestTransform(opt.datatype, True)), batch_size=opt.test_batch, num_workers=opt.workers, pin_memory=pin_memory) print('initializing model ...') if opt.model_name == 'softmax' or opt.model_name == 'softmax_triplet': model = ResNetBuilder(dataset.num_train_pids, 1, True) elif opt.model_name == 'triplet': model = ResNetBuilder(None, 1, True) elif opt.model_name == 'bfe': if opt.datatype == "person": model = BFE(dataset.num_train_pids, 1.0, 0.33) else: model = BFE(dataset.num_train_pids, 0.5, 0.5) elif opt.model_name == 'ide': model = IDE(dataset.num_train_pids) elif opt.model_name == 'resnet': model = Resnet(dataset.num_train_pids) elif opt.model_name == 'strongBaseline': model = StrongBaseline(dataset.num_train_pids) optim_policy = model.get_optim_policy() # update model model = resnet18(True) model.fc.out_features = 10 print('model size: {:.5f}M'.format( sum(p.numel() for p in model.parameters()) / 1e6)) if use_gpu: model = nn.DataParallel(model).cuda() # get optimizer if opt.optim == "sgd": optimizer = torch.optim.SGD(optim_policy, lr=opt.lr, momentum=0.9, weight_decay=opt.weight_decay) else: optimizer = torch.optim.Adam(optim_policy, lr=opt.lr, weight_decay=opt.weight_decay) #xent_criterion = nn.CrossEntropyLoss() criterion = CrossEntropyLabelSmooth(10) epochs = 100 best = 0.0 b_e = 0 for e in range(epochs): model.train() for i, inputs in enumerate(trainloader): imgs, pid, _ = inputs imgs, pid = imgs.cuda(), pid.cuda() outputs = model(imgs) loss = criterion(outputs, pid) optimizer.zero_grad() loss.backward() optimizer.step() print(('epoch=%s \t batch loss=%s') % (e, loss.item()))