示例#1
0
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))
示例#2
0
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))
示例#3
0
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))
示例#5
0
        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)]+\
示例#6
0
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()))