Relabel maps generating code about relabel_imagenet HOT 1 OPEN

normalize = torchvision.transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)
train_root = os.path.join(data_path, 'train')
train_dataset = ImageFolderWithPath(
    root=train_root,
    transform=torchvision.transforms.Compose([
        torchvision.transforms.Resize((machine_annotator_image_size,
                                       machine_annotator_image_size)),
        torchvision.transforms.ToTensor(),
        normalize,
    ]))

train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_sizebatch_size,
    shuffle=False,
    num_workers=machine_annotator_num_workers,
    pin_memory=True,
    sampler=train_sampler
)

score_maps = {}

print('Train loader size: {}'.format(len(train_loader)))
for iteration, batch in enumerate(train_loader):
    start_time = time.time()

    input, target, img_names = batch
    target = target.cuda()

    with torch.no_grad():
        # get score_map by removing GlobalAvergaPool. (See Appendix)
        output_map = machine_annotator_model(input, score_map=True) 

    output_topk = output_map.topk(k=machine_annotator_topk, dim=1)

    for idx, (t0, t1) in enumerate(zip(output_topk[0], output_topk[1])):
        img_name = img_names[idx].split('/')[-1].split('.')[0]
        t = torch.cat([t0.unsqueeze(0), t1.unsqueeze(0).float()],
                      dim=0)

        score_maps[img_name] = t.cpu()

score_map_path = '{}_{}_sz{}_top{}_score_maps.pt'.format(
    data.dataset,
    machine_annotator_arch,
    machine_annotator_image_size,
    machine_annotator_topk)