Ejemplo n.º 1
0
def main_worker(confg):
    global start_epoch, best_mAP

    cudnn.benchmark = True  #

    sys.stdout = Logger(os.path.join(config['logs_dir'], 'log.txt'))
    print("==========\nArgs:{}\n==========".format(config))

    #######################################################
    # create test data loaders
    #######################################################

    if config['iters'] > 0:
        iters = config['iters']
    else:
        iters = None

    #dataset_target = get_data(config['dataset_target'], config['data_dir'])
    test_loader_target = get_test_loader(config['data_dir'],
                                         config['resize_size'],
                                         config['batch_size'],
                                         config['workers'])

    #######################################################
    # create model
    #######################################################

    model_1, model_2, model_1_ema, model_2_ema = create_model(config)
Ejemplo n.º 2
0
def main_worker(args):
    cudnn.benchmark = True

    log_dir = osp.dirname(args.resume)
    sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    query_loader, gallery_loader = \
        get_data(args.dataset_target, args.data_dir, args.height,
                 args.width, args.batch_size, args.workers)

    # Create model
    model = models.create(args.arch, pretrained=False, num_features=args.features, dropout=args.dropout, num_classes=0)
    model.cuda()
    model = nn.DataParallel(model)

    # Load from checkpoint
    checkpoint = load_checkpoint(args.resume)
    copy_state_dict(checkpoint['state_dict'], model)
    start_epoch = checkpoint['epoch']
    best_mAP = checkpoint['best_mAP']
    print("=> Checkpoint of epoch {}  best mAP {:.1%}".format(start_epoch, best_mAP))

    # Evaluator
    q_f, q_id, q_cam = test(model, query_loader)
    g_f, g_id, g_cam = test(model, gallery_loader)

    q_g_dist = np.dot(q_f, np.transpose(g_f))
    q_g_dist = 2. - 2 * q_g_dist  # change the cosine similarity metric to euclidean similarity metric
    all_cmc, mAP = eval_func(q_g_dist, q_id, g_id, q_cam, g_cam)
    all_cmc = all_cmc * 100
    print('rank-1: {:.4f} rank-5: {:.4f} rank-10: {:.4f} rank-20: {:.4f} rank-50: {:.4f} mAP: {:.4f}'.format(all_cmc[0],
                                                                                                             all_cmc[4],
                                                                                                             all_cmc[9],
                                                                                                             all_cmc[
                                                                                                                 19],
                                                                                                             all_cmc[
                                                                                                                 49],
                                                                                                             mAP * 100))

    indices = np.argsort(q_g_dist, axis=1)
    np.savetxt("answer.txt", indices[:, :100], fmt="%04d")
    return
Ejemplo n.º 3
0
def main_worker(args):
    cudnn.benchmark = True

    log_dir = osp.dirname(args.resume)
    sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    dataset_target, test_loader_target = \
        get_data(args.dataset_target, args.data_dir, args.height,
                 args.width, args.batch_size, args.workers)

    # Create model
    model = models.create(args.arch,
                          pretrained=False,
                          num_features=args.features,
                          dropout=args.dropout,
                          num_classes=0)
    model.cuda()
    model = nn.DataParallel(model)

    # Load from checkpoint
    checkpoint = load_checkpoint(args.resume)
    copy_state_dict(checkpoint['state_dict'], model)
    start_epoch = checkpoint['epoch']
    best_mAP = checkpoint['best_mAP']
    print("=> Checkpoint of epoch {}  best mAP {:.1%}".format(
        start_epoch, best_mAP))

    # Evaluator
    evaluator = Evaluator(model)
    print("Test on the target domain of {}:".format(args.dataset_target))
    evaluator.evaluate(test_loader_target,
                       dataset_target.query,
                       dataset_target.gallery,
                       cmc_flag=True,
                       rerank=args.rerank)
    return
Ejemplo n.º 4
0
def main_worker(args):
    global best_mAP

    cudnn.benchmark = True

    sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    iters = args.iters if (args.iters > 0) else None
    dataset_target = get_data(args.dataset_target, args.data_dir)
    test_loader_target = get_test_loader(dataset_target, args.height,
                                         args.width, args.batch_size,
                                         args.workers)
    cluster_loader = get_test_loader(dataset_target,
                                     args.height,
                                     args.width,
                                     args.batch_size,
                                     args.workers,
                                     testset=dataset_target.train)

    # Create model
    model = create_model(args)

    # Evaluator
    evaluator = Evaluator(model)

    for epoch in range(args.epochs):
        dict_f, _ = extract_features(model, cluster_loader, print_freq=50)
        cf = torch.stack(list(dict_f.values())).numpy()

        print('\n Clustering into {} classes \n'.format(args.num_clusters))
        km = KMeans(n_clusters=args.num_clusters,
                    random_state=args.seed,
                    n_jobs=2).fit(cf)

        model.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_,
                                       axis=1)).float().cuda())

        target_label = km.labels_

        # change pseudo labels
        for i in range(len(dataset_target.train)):
            dataset_target.train[i] = list(dataset_target.train[i])
            dataset_target.train[i][1] = int(target_label[i])
            dataset_target.train[i] = tuple(dataset_target.train[i])

        train_loader_target = get_train_loader(dataset_target, args.height,
                                               args.width, args.batch_size,
                                               args.workers,
                                               args.num_instances, iters)

        # Optimizer
        params = []
        for key, value in model.named_parameters():
            if not value.requires_grad:
                continue
            params += [{
                "params": [value],
                "lr": args.lr,
                "weight_decay": args.weight_decay
            }]
        optimizer = torch.optim.Adam(params)

        # Trainer
        trainer = ClusterBaseTrainer(model, num_cluster=args.num_clusters)

        train_loader_target.new_epoch()

        trainer.train(epoch,
                      train_loader_target,
                      optimizer,
                      print_freq=args.print_freq,
                      train_iters=len(train_loader_target))

        def save_model(model, is_best, best_mAP):
            save_checkpoint(
                {
                    'state_dict': model.state_dict(),
                    'epoch': epoch + 1,
                    'best_mAP': best_mAP,
                },
                is_best,
                fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))

        if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1):
            mAP = evaluator.evaluate(test_loader_target,
                                     dataset_target.query,
                                     dataset_target.gallery,
                                     cmc_flag=False)
            is_best = (mAP > best_mAP)
            best_mAP = max(mAP, best_mAP)
            save_model(model, is_best, best_mAP)

            print(
                '\n * Finished epoch {:3d}  model mAP: {:5.1%} best: {:5.1%}{}\n'
                .format(epoch, mAP, best_mAP, ' *' if is_best else ''))

    print('Test on the best model.')
    checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
    model.load_state_dict(checkpoint['state_dict'])
    evaluator.evaluate(test_loader_target,
                       dataset_target.query,
                       dataset_target.gallery,
                       cmc_flag=True)
Ejemplo n.º 5
0
def main_worker(args):
    global start_epoch, best_mAP

    cudnn.benchmark = True

    if not args.evaluate:
        sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    else:
        log_dir = osp.dirname(args.resume)
        sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    iters = args.iters if (args.iters > 0) else None
    dataset_source, num_classes, train_loader_source = \
        get_data(args.dataset_source, args.data_dir, args.height,
                 args.width, args.batch_size, args.workers, args.num_instances, iters)
    dataset_target, _, train_loader_target = \
        get_data(args.dataset_target, args.data_dir, args.height,
                 args.width, args.batch_size, args.workers, args.num_instances, iters)
    dataset_validation, test_loader_target = \
        get_test_data(args.dataset_validation, args.data_dir, args.height,
                      args.width, args.batch_size, args.workers)

    # Create model
    model = models.create(args.arch,
                          dropout=args.dropout,
                          num_classes=num_classes,
                          circle=args.circle)
    # print(model)
    print("Model size: {:.3f} M".format(count_num_param(model)))
    model.cuda()
    model = nn.DataParallel(model)  # 多gpu并行

    # Load from checkpoint
    if args.resume:
        checkpoint = load_checkpoint(args.resume)
        copy_state_dict(checkpoint['state_dict'], model)
        start_epoch = checkpoint['epoch']
        best_mAP = checkpoint['best_mAP']
        print("=> Start epoch {}  best mAP {:.1%}".format(
            start_epoch, best_mAP))

    # Evaluator
    evaluator = Evaluator(model)
    if args.evaluate:
        # print("Test on source domain:")
        # evaluator.evaluate(test_loader_source, dataset_source.query, dataset_source.gallery, cmc_flag=True,
        #                    rerank=args.rerank)
        print("Test on target domain:")
        evaluator.evaluate(test_loader_target,
                           dataset_target.query,
                           dataset_target.gallery,
                           cmc_flag=True,
                           rerank=args.rerank)
        return

    params = []
    for key, value in model.named_parameters():
        if not value.requires_grad:
            continue
        params += [{
            "params": [value],
            "lr": args.lr,
            "weight_decay": args.weight_decay
        }]
    optimizer = torch.optim.Adam(params)
    lr_scheduler = WarmupMultiStepLR(optimizer,
                                     args.milestones,
                                     gamma=0.1,
                                     warmup_factor=0.01,
                                     warmup_iters=args.warmup_step)

    # Trainer
    trainer = PreTrainer(model, num_classes, args, margin=args.margin)

    # Start training
    for epoch in range(start_epoch, args.epochs):
        lr_scheduler.step()
        train_loader_source.new_epoch()
        # train_loader_target.new_epoch()

        trainer.train(epoch,
                      train_loader_source,
                      optimizer,
                      train_iters=len(train_loader_source),
                      print_freq=args.print_freq,
                      balance=args.balance)

        if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1):
            _, mAP = evaluator.evaluate(test_loader_target,
                                        dataset_validation.query,
                                        dataset_validation.gallery,
                                        cmc_flag=True)

            is_best = mAP > best_mAP
            best_mAP = max(mAP, best_mAP)
            save_checkpoint(
                {
                    'state_dict': model.state_dict(),
                    'epoch': epoch + 1,
                    'best_mAP': best_mAP,
                },
                is_best,
                fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))

            print(
                '\n * Finished epoch {:3d}  source mAP: {:5.1%}  best: {:5.1%}{}\n'
                .format(epoch, mAP, best_mAP, ' *' if is_best else ''))

    print("Test on target domain:")
    evaluator.evaluate(test_loader_target,
                       dataset_validation.query,
                       dataset_validation.gallery,
                       cmc_flag=True,
                       rerank=args.rerank)
Ejemplo n.º 6
0
def main_worker(args):
    global start_epoch, best_mAP

    cudnn.benchmark = True

    sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    iters = args.iters if (args.iters > 0) else None
    dataset_source = get_data(args.dataset_source, args.data_dir)
    dataset_target = get_data(args.dataset_target, args.data_dir)
    test_loader_target = get_test_loader(dataset_target, args.height,
                                         args.width, args.batch_size,
                                         args.workers)
    tar_cluster_loader = get_test_loader(dataset_target,
                                         args.height,
                                         args.width,
                                         args.batch_size,
                                         args.workers,
                                         testset=dataset_target.train)
    sour_cluster_loader = get_test_loader(dataset_source,
                                          args.height,
                                          args.width,
                                          args.batch_size,
                                          args.workers,
                                          testset=dataset_source.train)

    # Create model
    model_1, model_2, model_1_ema, model_2_ema = create_model(
        args, len(dataset_target.train))

    # Evaluator
    evaluator_1_ema = Evaluator(model_1_ema)
    evaluator_2_ema = Evaluator(model_2_ema)

    for epoch in range(args.epochs):
        dict_f, _ = extract_features(model_1_ema,
                                     tar_cluster_loader,
                                     print_freq=50)
        cf_1 = torch.stack(list(dict_f.values()))
        dict_f, _ = extract_features(model_2_ema,
                                     tar_cluster_loader,
                                     print_freq=50)
        cf_2 = torch.stack(list(dict_f.values()))
        cf = (cf_1 + cf_2) / 2
        cf = F.normalize(cf, dim=1)

        if (args.lambda_value > 0):
            dict_f, _ = extract_features(model_1_ema,
                                         sour_cluster_loader,
                                         print_freq=50)
            cf_1 = torch.stack(list(dict_f.values()))
            dict_f, _ = extract_features(model_2_ema,
                                         sour_cluster_loader,
                                         print_freq=50)
            cf_2 = torch.stack(list(dict_f.values()))
            cf_s = (cf_1 + cf_2) / 2
            cf_s = F.normalize(cf_s, dim=1)
            rerank_dist = compute_jaccard_dist(cf,
                                               lambda_value=args.lambda_value,
                                               source_features=cf_s,
                                               use_gpu=args.rr_gpu).numpy()
        else:
            rerank_dist = compute_jaccard_dist(cf, use_gpu=args.rr_gpu).numpy()

        if (epoch == 0):
            # DBSCAN cluster
            tri_mat = np.triu(rerank_dist, 1)  # tri_mat.dim=2
            tri_mat = tri_mat[np.nonzero(tri_mat)]  # tri_mat.dim=1
            tri_mat = np.sort(tri_mat, axis=None)
            rho = 1.6e-3
            top_num = np.round(rho * tri_mat.size).astype(int)
            eps = tri_mat[:top_num].mean()
            print('eps for cluster: {:.3f}'.format(eps))
            cluster = DBSCAN(eps=eps,
                             min_samples=4,
                             metric='precomputed',
                             n_jobs=-1)

        print('Clustering and labeling...')
        labels = cluster.fit_predict(rerank_dist)
        num_ids = len(set(labels)) - (1 if -1 in labels else 0)
        args.num_clusters = num_ids
        print('\n Clustered into {} classes \n'.format(args.num_clusters))

        # generate new dataset and calculate cluster centers
        new_dataset = []
        cluster_centers = collections.defaultdict(list)
        for i, ((fname, _, cid),
                label) in enumerate(zip(dataset_target.train, labels)):
            if label == -1: continue
            new_dataset.append((fname, label, cid))
            cluster_centers[label].append(cf[i])

        cluster_centers = [
            torch.stack(cluster_centers[idx]).mean(0)
            for idx in sorted(cluster_centers.keys())
        ]
        cluster_centers = torch.stack(cluster_centers)
        model_1.module.classifier.weight.data[:args.num_clusters].copy_(
            F.normalize(cluster_centers, dim=1).float().cuda())
        model_2.module.classifier.weight.data[:args.num_clusters].copy_(
            F.normalize(cluster_centers, dim=1).float().cuda())
        model_1_ema.module.classifier.weight.data[:args.num_clusters].copy_(
            F.normalize(cluster_centers, dim=1).float().cuda())
        model_2_ema.module.classifier.weight.data[:args.num_clusters].copy_(
            F.normalize(cluster_centers, dim=1).float().cuda())

        train_loader_target = get_train_loader(dataset_target,
                                               args.height,
                                               args.width,
                                               args.batch_size,
                                               args.workers,
                                               args.num_instances,
                                               iters,
                                               trainset=new_dataset)

        # Optimizer
        params = []
        for key, value in model_1.named_parameters():
            if not value.requires_grad:
                continue
            params += [{
                "params": [value],
                "lr": args.lr,
                "weight_decay": args.weight_decay
            }]
        for key, value in model_2.named_parameters():
            if not value.requires_grad:
                continue
            params += [{
                "params": [value],
                "lr": args.lr,
                "weight_decay": args.weight_decay
            }]
        optimizer = torch.optim.Adam(params)

        # Trainer
        trainer = MMTTrainer(model_1,
                             model_2,
                             model_1_ema,
                             model_2_ema,
                             num_cluster=args.num_clusters,
                             alpha=args.alpha)

        train_loader_target.new_epoch()

        trainer.train(epoch,
                      train_loader_target,
                      optimizer,
                      ce_soft_weight=args.soft_ce_weight,
                      tri_soft_weight=args.soft_tri_weight,
                      print_freq=args.print_freq,
                      train_iters=len(train_loader_target))

        def save_model(model_ema, is_best, best_mAP, mid):
            save_checkpoint(
                {
                    'state_dict': model_ema.state_dict(),
                    'epoch': epoch + 1,
                    'best_mAP': best_mAP,
                },
                is_best,
                fpath=osp.join(args.logs_dir,
                               'model' + str(mid) + '_checkpoint.pth.tar'))

        if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)):
            mAP_1 = evaluator_1_ema.evaluate(test_loader_target,
                                             dataset_target.query,
                                             dataset_target.gallery,
                                             cmc_flag=False)
            mAP_2 = evaluator_2_ema.evaluate(test_loader_target,
                                             dataset_target.query,
                                             dataset_target.gallery,
                                             cmc_flag=False)
            is_best = (mAP_1 > best_mAP) or (mAP_2 > best_mAP)
            best_mAP = max(mAP_1, mAP_2, best_mAP)
            save_model(model_1_ema, (is_best and (mAP_1 > mAP_2)), best_mAP, 1)
            save_model(model_2_ema, (is_best and (mAP_1 <= mAP_2)), best_mAP,
                       2)

            print(
                '\n * Finished epoch {:3d}  model no.1 mAP: {:5.1%} model no.2 mAP: {:5.1%}  best: {:5.1%}{}\n'
                .format(epoch, mAP_1, mAP_2, best_mAP,
                        ' *' if is_best else ''))

    print('Test on the best model.')
    checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
    model_1_ema.load_state_dict(checkpoint['state_dict'])
    evaluator_1_ema.evaluate(test_loader_target,
                             dataset_target.query,
                             dataset_target.gallery,
                             cmc_flag=True)
Ejemplo n.º 7
0
def main_worker(args):
    global best_mAP

    cudnn.benchmark = True

    sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    iters = args.iters if (args.iters > 0) else None
    dataset_target = get_data(args.dataset_target, args.data_dir)
    test_loader_target = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers)
    # test loader target 包含data set target的query和gallery
    cluster_loader = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers,
                                     testset=dataset_target.train)
    # cluster loader 只包含data set target 的 train

    # Create model
    model_1, model_2, model_1_ema, model_2_ema = create_model(args)

    # Evaluator
    evaluator_1_ema = Evaluator(model_1_ema)
    evaluator_2_ema = Evaluator(model_2_ema)

    for epoch in range(args.epochs):
        dict_f, _ = extract_features(model_1_ema, cluster_loader, print_freq=50)  # 返回的是features labels的字典
        cf_1 = torch.stack(list(dict_f.values())).numpy()
        # dict_f.values()返回字典中的所有值
        # list()将元组转换为列表
        # torch.stack()沿着一个新维度对输入张量序列进行扩维连接,dim参数默认为0,形成n*x*y张量
        dict_f, _ = extract_features(model_2_ema, cluster_loader, print_freq=50)
        cf_2 = torch.stack(list(dict_f.values())).numpy()
        cf = (cf_1 + cf_2) / 2

        print('\n Clustering into {} classes \n'.format(args.num_clusters))
        km = KMeans(n_clusters=args.num_clusters, random_state=args.seed, n_jobs=-1).fit(cf)

        # 聚类后的各个类的中心,将这些中心标准化以后直接赋值到了线性层的参数中
        # 我们假设一个线性层的输入特征是M维的,输出是N维的,如果不加思考的化,就是一个M维到N维的转换。但是深入一点的思考的化,其实
        # 可以将其理解为:将M维的向量,与N个M维的向量(Normalize后则是单位向量)做内积,生成了N个内积,内积i代表着输入特征在第i个
        # 向量上的投影,这个值越大,代表着输入向量与这个向量的关系越紧密。
        model_1.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda())
        model_2.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda())
        model_1_ema.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda())
        model_2_ema.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda())

        target_label = km.labels_

        # change pseudo labels
        for i in range(len(dataset_target.train)):
            dataset_target.train[i] = list(dataset_target.train[i])
            dataset_target.train[i][1] = int(target_label[i])
            dataset_target.train[i] = tuple(dataset_target.train[i])

        train_loader_target = get_train_loader(dataset_target, args.height, args.width,
                                               args.batch_size, args.workers, args.num_instances, iters)

        # Optimizer
        params = []
        for key, value in model_1.named_parameters():
            if not value.requires_grad:
                continue
            params += [{"params": [value], "lr": args.lr, "weight_decay": args.weight_decay}]
        for key, value in model_2.named_parameters():
            if not value.requires_grad:
                continue
            params += [{"params": [value], "lr": args.lr, "weight_decay": args.weight_decay}]
        optimizer = torch.optim.Adam(params)

        # Trainer
        trainer = MMTTrainer(model_1, model_2, model_1_ema, model_2_ema, args)

        train_loader_target.new_epoch()

        trainer.train(epoch, train_loader_target, optimizer,
                      ce_soft_weight=args.soft_ce_weight, tri_soft_weight=args.soft_tri_weight,
                      print_freq=args.print_freq, train_iters=len(train_loader_target), balance=args.balance)

        def save_model(model_ema, is_best, best_mAP, mid):
            save_checkpoint({
                'state_dict': model_ema.state_dict(),
                'epoch': epoch + 1,
                'best_mAP': best_mAP,
            }, is_best, fpath=osp.join(args.logs_dir, 'model' + str(mid) + '_checkpoint.pth.tar'))

        if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1):
            mAP_1 = evaluator_1_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery,
                                             cmc_flag=False)
            mAP_2 = evaluator_2_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery,
                                             cmc_flag=False)
            is_best = (mAP_1 > best_mAP) or (mAP_2 > best_mAP)
            best_mAP = max(mAP_1, mAP_2, best_mAP)
            save_model(model_1_ema, (is_best and (mAP_1 > mAP_2)), best_mAP, 1)
            save_model(model_2_ema, (is_best and (mAP_1 <= mAP_2)), best_mAP, 2)

            print('\n * Finished epoch {:3d}  model no.1 mAP: {:5.1%} model no.2 mAP: {:5.1%}  best: {:5.1%}{}\n'.
                  format(epoch, mAP_1, mAP_2, best_mAP, ' *' if is_best else ''))

    print('Test on the best model.')
    checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
    model_1_ema.load_state_dict(checkpoint['state_dict'])
    evaluator_1_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True)
Ejemplo n.º 8
0
def main_worker(args):
    global start_epoch, best_mAP

    cudnn.benchmark = True

    sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    # Create data loaders
    iters = args.iters if (args.iters > 0) else None
    dataset_target = get_data(args.dataset_target, args.data_dir)
    test_loader_target = get_test_loader(dataset_target, args.height,
                                         args.width, args.batch_size,
                                         args.workers)

    # Create model
    model_1, model_2, model_1_ema, model_2_ema = create_model(args)

    # Evaluator
    evaluator_1_ema = Evaluator(model_1_ema)
    evaluator_2_ema = Evaluator(model_2_ema)

    clusters = [args.num_clusters] * args.epochs
    feature_length = args.features if args.features > 0 else 2048
    moving_avg_features = np.zeros((len(dataset_target.train), feature_length))

    for nc in range(len(clusters)):
        cluster_loader = get_test_loader(dataset_target,
                                         args.height,
                                         args.width,
                                         args.batch_size,
                                         args.workers,
                                         testset=dataset_target.train)

        dict_f, _ = extract_features(model_1_ema,
                                     cluster_loader,
                                     print_freq=50)
        cf_1 = torch.stack(list(dict_f.values())).numpy()
        dict_f, _ = extract_features(model_2_ema,
                                     cluster_loader,
                                     print_freq=50)
        cf_2 = torch.stack(list(dict_f.values())).numpy()
        cf = (cf_1 + cf_2) / 2

        moving_avg_features = moving_avg_features * args.moving_avg_momentum + cf * (
            1 - args.moving_avg_momentum)
        moving_avg_features = moving_avg_features / (
            1 - args.moving_avg_momentum**(nc + 1))

        print('\n Clustering into {} classes \n'.format(clusters[nc]))
        km = KMeans(n_clusters=clusters[nc], random_state=args.seed,
                    n_jobs=2).fit(moving_avg_features)

        model_1.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_,
                                       axis=1)).float().cuda())
        model_2.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_,
                                       axis=1)).float().cuda())
        model_1_ema.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_,
                                       axis=1)).float().cuda())
        model_2_ema.module.classifier.weight.data.copy_(
            torch.from_numpy(normalize(km.cluster_centers_,
                                       axis=1)).float().cuda())

        target_label = km.labels_

        # change pseudo labels
        for i in range(len(dataset_target.train)):
            dataset_target.train[i] = list(dataset_target.train[i])
            dataset_target.train[i][1] = int(target_label[i])
            dataset_target.train[i] = tuple(dataset_target.train[i])

        train_loader_target = get_train_loader(dataset_target, args.height,
                                               args.width, args.batch_size,
                                               args.workers,
                                               args.num_instances, iters)

        # Optimizer
        params = []
        for key, value in model_1.named_parameters():
            if not value.requires_grad:
                continue
            params += [{
                "params": [value],
                "lr": args.lr,
                "weight_decay": args.weight_decay
            }]
        for key, value in model_2.named_parameters():
            if not value.requires_grad:
                continue
            params += [{
                "params": [value],
                "lr": args.lr,
                "weight_decay": args.weight_decay
            }]
        optimizer = torch.optim.Adam(params)

        # Trainer
        trainer = MMTTrainer(model_1,
                             model_2,
                             model_1_ema,
                             model_2_ema,
                             num_cluster=clusters[nc],
                             alpha=args.alpha)

        train_loader_target.new_epoch()
        epoch = nc

        trainer.train(epoch,
                      train_loader_target,
                      optimizer,
                      ce_soft_weight=args.soft_ce_weight,
                      tri_soft_weight=args.soft_tri_weight,
                      print_freq=args.print_freq,
                      train_iters=len(train_loader_target))

        def save_model(model_ema, is_best, best_mAP, mid):
            save_checkpoint(
                {
                    'state_dict': model_ema.state_dict(),
                    'epoch': epoch + 1,
                    'best_mAP': best_mAP,
                },
                is_best,
                fpath=osp.join(args.logs_dir,
                               'model' + str(mid) + '_checkpoint.pth.tar'))

        if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)):
            mAP_1 = evaluator_1_ema.evaluate(test_loader_target,
                                             dataset_target.query,
                                             dataset_target.gallery,
                                             cmc_flag=False)
            mAP_2 = evaluator_2_ema.evaluate(test_loader_target,
                                             dataset_target.query,
                                             dataset_target.gallery,
                                             cmc_flag=False)
            is_best = (mAP_1 > best_mAP) or (mAP_2 > best_mAP)
            best_mAP = max(mAP_1, mAP_2, best_mAP)
            save_model(model_1_ema, (is_best and (mAP_1 > mAP_2)), best_mAP, 1)
            save_model(model_2_ema, (is_best and (mAP_1 <= mAP_2)), best_mAP,
                       2)

            print(
                '\n * Finished epoch {:3d}  model no.1 mAP: {:5.1%} model no.2 mAP: {:5.1%}  best: {:5.1%}{}\n'
                .format(epoch, mAP_1, mAP_2, best_mAP,
                        ' *' if is_best else ''))

    print('Test on the best model.')
    checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
    model_1_ema.load_state_dict(checkpoint['state_dict'])
    evaluator_1_ema.evaluate(test_loader_target,
                             dataset_target.query,
                             dataset_target.gallery,
                             cmc_flag=True)