def main(rank, args):
    dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    seed = (args.seed + dist.get_rank())
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.cuda.set_device(rank)
    trainset = DataFactory(name=args.dataset, partition=args.partitions[0], data_root=args.data_root)
    testset = DataFactory(name=args.dataset, partition=args.partitions[1], data_root=args.data_root)
    train_loader = DataLoader(dataset=trainset, collate_fn=custom_collate, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, drop_last=True, sampler=DistributedSampler(trainset, num_replicas=args.world_size, rank=rank))
    test_loader = DataLoader(dataset=testset, collate_fn=custom_collate, batch_size=1, num_workers=args.num_workers, pin_memory=True, drop_last=False, sampler=torch.utils.data.SequentialSampler(testset))
    args.human_idx = 0
    if (args.dataset == 'hicodet'):
        object_to_target = train_loader.dataset.dataset.object_to_verb
        args.num_classes = 117
    elif (args.dataset == 'vcoco'):
        object_to_target = list(train_loader.dataset.dataset.object_to_action.values())
        args.num_classes = 24
    upt = build_detector(args, object_to_target)
    if os.path.exists(args.resume):
        print(f'=> Rank {rank}: continue from saved checkpoint {args.resume}')
        checkpoint = torch.load(args.resume, map_location='cpu')
        upt.load_state_dict(checkpoint['model_state_dict'])
    else:
        print(f'=> Rank {rank}: start from a randomly initialised model')
    engine = CustomisedDLE(upt, train_loader, max_norm=args.clip_max_norm, num_classes=args.num_classes, print_interval=args.print_interval, find_unused_parameters=True, cache_dir=args.output_dir)
    if args.cache:
        if (args.dataset == 'hicodet'):
            engine.cache_hico(test_loader, args.output_dir)
        elif (args.dataset == 'vcoco'):
            engine.cache_vcoco(test_loader, args.output_dir)
        return
    if args.eval:
        if (args.dataset == 'vcoco'):
            raise NotImplementedError(f'Evaluation on V-COCO has not been implemented.')
        ap = engine.test_hico(test_loader)
        num_anno = torch.as_tensor(trainset.dataset.anno_interaction)
        rare = torch.nonzero((num_anno < 10)).squeeze(1)
        non_rare = torch.nonzero((num_anno >= 10)).squeeze(1)
        print(f'The mAP is {ap.mean():.4f}, rare: {ap[rare].mean():.4f}, none-rare: {ap[non_rare].mean():.4f}')
        return
    for p in upt.detector.parameters():
        p.requires_grad = False
    param_dicts = [{'params': [p for (n, p) in upt.named_parameters() if (('interaction_head' in n) and p.requires_grad)]}]
    optim = torch.optim.AdamW(param_dicts, lr=args.lr_head, weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, args.lr_drop)
    engine.update_state_key(optimizer=optim, lr_scheduler=lr_scheduler)
    engine(args.epochs)
def main(rank, args):
    dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    seed = (args.seed + dist.get_rank())
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.cuda.set_device(rank)
    trainset = DataFactory(name=args.dataset, partition=args.partitions[0], data_root=args.data_root)
    testset = DataFactory(name=args.dataset, partition=args.partitions[1], data_root=args.data_root)
    train_loader = DataLoader(dataset=trainset, collate_fn=custom_collate, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True, drop_last=True, sampler=DistributedSampler(trainset, num_replicas=args.world_size, rank=rank))
    test_loader = DataLoader(dataset=testset, collate_fn=custom_collate, batch_size=1, num_workers=args.num_workers, pin_memory=True, drop_last=False, sampler=torch.utils.data.SequentialSampler(testset))
    args.human_idx = 0
    if (args.dataset == 'hicodet'):
        object_to_target = train_loader.dataset.dataset.object_to_verb
        args.num_classes = 117
    elif (args.dataset == 'vcoco'):
        object_to_target = list(train_loader.dataset.dataset.object_to_action.values())
        args.num_classes = 24
    upt = build_detector(args, object_to_target)
    if os.path.exists(args.resume):
        print(f'=> Rank {rank}: continue from saved checkpoint {args.resume}')
        checkpoint = torch.load(args.resume, map_location='cpu')
        upt.load_state_dict(checkpoint['model_state_dict'])
    else:
        print(f'=> Rank {rank}: start from a randomly initialised model')
    engine = CustomisedDLE(upt, train_loader, max_norm=args.clip_max_norm, num_classes=args.num_classes, print_interval=args.print_interval, find_unused_parameters=True, cache_dir=args.output_dir)
    if args.cache:
        if (args.dataset == 'hicodet'):
            engine.cache_hico(test_loader, args.output_dir)
        elif (args.dataset == 'vcoco'):
            engine.cache_vcoco(test_loader, args.output_dir)
        return
    if args.eval:
        if (args.dataset == 'vcoco'):
            raise NotImplementedError(f'Evaluation on V-COCO has not been implemented.')
        ap = engine.test_hico(test_loader)
        num_anno = torch.as_tensor(trainset.dataset.anno_interaction)
        rare = torch.nonzero((num_anno < 10)).squeeze(1)
        non_rare = torch.nonzero((num_anno >= 10)).squeeze(1)
        print(f'The mAP is {ap.mean():.4f}, rare: {ap[rare].mean():.4f}, none-rare: {ap[non_rare].mean():.4f}')
        return
    for p in upt.detector.parameters():
        p.requires_grad = False
    param_dicts = [{'params': [p for (n, p) in upt.named_parameters() if (('interaction_head' in n) and p.requires_grad)]}]
    optim = torch.optim.AdamW(param_dicts, lr=args.lr_head, weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, args.lr_drop)
    engine.update_state_key(optimizer=optim, lr_scheduler=lr_scheduler)
    engine(args.epochs)