def main_worker(args):
    global start_epoch, best_mAP
    start_time = time.monotonic()
    cudnn.benchmark = True
    sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    print('==========\nArgs:{}\n=========='.format(args))
    iters = (args.iters if (args.iters > 0) else None)
    print('==> Load source-domain dataset')
    train_items = []
    for src in args.dataset_source.split(','):
        dataset = get_data(src, args.data_dir, args.combine_all)
        train_items.extend(dataset.train)
    dataset_source = CommDataset(train_items)
    print('==> Load target-domain dataset')
    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)
    train_loader_source = get_train_loader(args, dataset_source, args.height, args.width, args.batch_size, args.workers, args.num_instances, iters)
    source_classes = dataset_source.num_train_pids
    args.nclass = source_classes
    model = create_model(args)
    print(model)
    evaluator = Evaluator(model)
    params = [{'params': [value]} for (_, value) in model.named_parameters() if value.requires_grad]
    optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=0.1)
    trainer = RSCTrainer(model, args.nclass, margin=args.margin)
    for epoch in range(args.epochs):
        train_loader_source.new_epoch()
        trainer.train(epoch, train_loader_source, optimizer, print_freq=args.print_freq, train_iters=args.iters)
        if ((((epoch + 1) % args.eval_step) == 0) or (epoch == (args.epochs - 1))):
            print('Test on target: ', args.dataset_target)
            (_, mAP) = evaluator.evaluate(test_loader_target, dataset_target.query, dataset_target.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}  model mAP: {:5.1%}  best: {:5.1%}{}\n'.format(epoch, mAP, best_mAP, (' *' if is_best else '')))
        lr_scheduler.step()
    print('==> Test with the best model on the target domain:')
    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)
    end_time = time.monotonic()
    print('Total running time: ', timedelta(seconds=(end_time - start_time)))
def main_worker(args):
    global start_epoch, best_mAP
    start_time = time.monotonic()
    cudnn.benchmark = True
    sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
    print('==========\nArgs:{}\n=========='.format(args))
    iters = (args.iters if (args.iters > 0) else None)
    print('==> Load source-domain dataset')
    train_items = []
    for src in args.dataset_source.split(','):
        dataset = get_data(src, args.data_dir, args.combine_all)
        train_items.extend(dataset.train)
    dataset_source = CommDataset(train_items)
    print('==> Load target-domain dataset')
    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)
    train_loader_source = get_train_loader(args, dataset_source, args.height, args.width, args.batch_size, args.workers, args.num_instances, iters)
    source_classes = dataset_source.num_train_pids
    args.nclass = source_classes
    model = create_model(args)
    print(model)
    evaluator = Evaluator(model)
    params = [{'params': [value]} for (_, value) in model.named_parameters() if value.requires_grad]
    optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=0.1)
    trainer = RSCTrainer(model, args.nclass, margin=args.margin)
    for epoch in range(args.epochs):
        train_loader_source.new_epoch()
        trainer.train(epoch, train_loader_source, optimizer, print_freq=args.print_freq, train_iters=args.iters)
        if ((((epoch + 1) % args.eval_step) == 0) or (epoch == (args.epochs - 1))):
            print('Test on target: ', args.dataset_target)
            (_, mAP) = evaluator.evaluate(test_loader_target, dataset_target.query, dataset_target.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}  model mAP: {:5.1%}  best: {:5.1%}{}\n'.format(epoch, mAP, best_mAP, (' *' if is_best else )))
        lr_scheduler.step()
    print('==> Test with the best model on the target domain:')
    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)
    end_time = time.monotonic()
    print('Total running time: ', timedelta(seconds=(end_time - start_time)))