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'))
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, test_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, test_loader_target = \
        get_data(args.dataset_target, args.data_dir, args.height,
                 args.width, args.batch_size, args.workers, 0, iters)

    # Create model
    model = models.create(args.arch,
                          num_features=args.features,
                          dropout=args.dropout,
                          num_classes=[num_classes])
    model.cuda()
    model = nn.DataParallel(model)
    print(model)
    # 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)
    # args.evaluate=True
    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, 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,
                      train_loader_target,
                      optimizer,
                      train_iters=len(train_loader_source),
                      print_freq=args.print_freq)

        if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)):

            _, mAP = evaluator.evaluate(test_loader_source,
                                        dataset_source.query,
                                        dataset_source.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_target.query,
                       dataset_target.gallery,
                       cmc_flag=True,
                       rerank=args.rerank)