update to torch2

app.py

@@ -0,0 +1,35 @@
+from inference import Inference
+import argparse
+import gradio as gr
+import glob
+
+def parse_option():
+    parser = argparse.ArgumentParser('MetaFG Inference script', add_help=False)
+    parser.add_argument('--cfg', type=str, metavar="FILE", help='path to config file', default="configs/MetaFG_2_224.yaml")
+    # easy config modification
+    parser.add_argument('--model-path', type=str, help="path to model data", default="./ckpt_4_mf2.pth")
+    parser.add_argument('--img-size', type=int, default=384, help='path to image')
+    parser.add_argument('--meta-path', default="meta.txt", type=str, help='path to meta data')
+    parser.add_argument('--names-path', default="names_mf2.txt", type=str, help='path to meta data')
+    args = parser.parse_args()
+    return args
+
+if __name__ == '__main__':
+    args = parse_option()
+
+    model = Inference(config_path=args.cfg,
+                       model_path=args.model_path,
+                       names_path=args.names_path)
+    
+    def classify(image):
+        preds = model.infer(img_path=image, meta_data_path="meta.txt").squeeze()
+        print(len(model.classes))
+        print(model.classes)
+        confidences = {c: float(preds[i]) for i,c in enumerate(model.classes)}
+
+        return confidences
+    
+    gr.Interface(pfn=classify, 
+            inputs=gr.Image(shape=(args.img_size, args.img_size), type="pil"),
+            outputs=gr.Label(num_top_classes=10),
+            examples=glob.glob("./example_images/*")).launch()

config.py

@@ -24,6 +24,8 @@ _C.DATA.BATCH_SIZE = 32
 _C.DATA.DATA_PATH = ''
 # Dataset name
 _C.DATA.DATASET = 'imagenet'
+# Dataset root folder
+_C.DATA.DATASET_ROOT = None
 # Input image size
 _C.DATA.IMG_SIZE = 224
 # Interpolation to resize image (random, bilinear, bicubic)
@@ -74,6 +76,7 @@ _C.MODEL.LABEL_SMOOTHING = 0.1
 _C.MODEL.PRETRAINED = None
 _C.MODEL.DORP_HEAD = True
 _C.MODEL.DORP_META = True
+_C.MODEL.FREEZE_BACKBONE = True
 
 _C.MODEL.ONLY_LAST_CLS = False
 _C.MODEL.EXTRA_TOKEN_NUM = 1
@@ -255,7 +258,7 @@ def update_config(config, args):
         config.MODEL.PRETRAINED = args.pretrain
 
     # set local rank for distributed training
-    config.LOCAL_RANK = args.local_rank
+    config.LOCAL_RANK = os.environ['LOCAL_RANK']
 
     # output folder
     config.OUTPUT = os.path.join(config.OUTPUT, config.MODEL.NAME, config.TAG)

data/build.py

@@ -13,7 +13,7 @@ from torchvision import datasets, transforms
 from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from timm.data import Mixup
 from timm.data import create_transform
-from timm.data.transforms import _pil_interp
+from timm.data.transforms import str_to_interp_mode
 
 from .cached_image_folder import CachedImageFolder
 from .samplers import SubsetRandomSampler
@@ -81,50 +81,41 @@ def build_dataset(is_train, config):
 #             root = os.path.join(config.DATA.DATA_PATH, prefix)
             root = './datasets/imagenet'
             dataset = datasets.ImageFolder(root, transform=transform)
-        nb_classes = 1000
     elif config.DATA.DATASET == 'inaturelist2021':
         root = './datasets/inaturelist2021'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 10000
     elif config.DATA.DATASET == 'inaturelist2021_mini':
         root = './datasets/inaturelist2021_mini'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 10000
     elif config.DATA.DATASET == 'inaturelist2017':
         root = './datasets/inaturelist2017'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 5089
     elif config.DATA.DATASET == 'inaturelist2018':
         root = './datasets/inaturelist2018'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 8142
     elif config.DATA.DATASET == 'cub-200':
         root = './datasets/cub-200'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 200
     elif config.DATA.DATASET == 'stanfordcars':
         root = './datasets/stanfordcars'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 196
     elif config.DATA.DATASET == 'oxfordflower':
         root = './datasets/oxfordflower'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 102
     elif config.DATA.DATASET == 'stanforddogs':
         root = './datasets/stanforddogs'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 120
     elif config.DATA.DATASET == 'nabirds':
         root = './datasets/nabirds'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 555
     elif config.DATA.DATASET == 'aircraft':
         root = './datasets/aircraft'
         dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
-        nb_classes = 100
     else:
-        raise NotImplementedError("We only support ImageNet and inaturelist.")
+        root = config.DATA.DATASET_ROOT
+        dataset = DatasetMeta(root=root,transform=transform,train=is_train,aux_info=config.DATA.ADD_META,dataset=config.DATA.DATASET)
 
+    nb_classes = len(dataset.class_to_idx)
     return dataset, nb_classes
 
 
@@ -153,14 +144,14 @@ def build_transform(is_train, config):
         if config.TEST.CROP:
             size = int((256 / 224) * config.DATA.IMG_SIZE)
             t.append(
-                transforms.Resize(size, interpolation=_pil_interp(config.DATA.INTERPOLATION)),
+                transforms.Resize(size, interpolation=str_to_interp_mode(config.DATA.INTERPOLATION)),
                 # to maintain same ratio w.r.t. 224 images
             )
             t.append(transforms.CenterCrop(config.DATA.IMG_SIZE))
         else:
             t.append(
                 transforms.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
-                                  interpolation=_pil_interp(config.DATA.INTERPOLATION))
+                                  interpolation=str_to_interp_mode(config.DATA.INTERPOLATION))
             )
 
     t.append(transforms.ToTensor())

data/dataset_fg.py

@@ -10,6 +10,7 @@ import pickle
 import numpy as np
 import pandas as pd
 import random
+from tqdm import tqdm
 random.seed(2021)
 from PIL import Image
 from scipy import io as scio
@@ -335,7 +336,7 @@ def find_images_and_targets_2017_2018(root,dataset,istrain=False,aux_info=False)
         else:
             images_and_targets.append((file_path,target))
     return images_and_targets,class_to_idx,images_info
-def find_images_and_targets(root,istrain=False,aux_info=False):
+def find_images_and_targets(root,istrain=False,aux_info=False, integrity_check=False):
     if os.path.exists(os.path.join(root,'train.json')):
         with open(os.path.join(root,'train.json'),'r') as f:
             train_class_info = json.load(f)
@@ -343,24 +344,59 @@ def find_images_and_targets(root,istrain=False,aux_info=False):
         with open(os.path.join(root,'train_mini.json'),'r') as f:
             train_class_info = json.load(f)
     else:
-        raise ValueError(f'not eixst file {root}/train.json or {root}/train_mini.json')
+        raise ValueError(f'{root}/train.json or {root}/train_mini.json doesn\'t exist')
+    
     with open(os.path.join(root,'val.json'),'r') as f:
         val_class_info = json.load(f)
-    categories_2021 = [x['name'].strip().lower() for x in val_class_info['categories']]
-    class_to_idx = {c: idx for idx, c in enumerate(categories_2021)}
+
+    categories = [x['name'].strip().lower() for x in val_class_info['categories']]
+    class_to_idx = {c: idx for idx, c in enumerate(categories)}
     id2label = dict()
     for categorie in train_class_info['categories']:
         id2label[int(categorie['id'])] = categorie['name'].strip().lower()
     class_info = train_class_info if istrain else val_class_info
-    
+    image_subdir = "train" if istrain else "val"
     images_and_targets = []
     images_info = []
     if aux_info:
         temporal_info = []
         spatial_info = []
 
-    for image,annotation in zip(class_info['images'],class_info['annotations']):
-        file_path = os.path.join(root,image['file_name'])
+    ann2im = {}
+    for ann in class_info['annotations']:
+        ann2im[ann['id']] = ann['image_id']
+    
+    ims = {}
+    for image in class_info['images']:
+        ims[image['id']] = image
+
+    print("Found", len(train_class_info['categories']))
+    print("Loading images and targets, checking image integrity")
+
+    for annotation in tqdm(class_info['annotations']):
+
+        image = ims[annotation['image_id']]
+        dir = train_class_info['categories'][annotation['category_id']]['image_dir_name']
+
+        file_path = os.path.join(root,image_subdir,dir,image['file_name'])
+
+        if not os.path.exists(file_path):
+
+            continue 
+        
+            print(f"Download {file_path}")
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            import requests
+            with open(file_path, 'wb') as fp:
+                fp.write(requests.get(image['inaturalist_url']).content)
+
+        if integrity_check:
+            try:
+                _ = np.array(Image.open(file_path))
+            except:
+                print(f"Failed to open {file_path}")
+                continue
+
         id_name = id2label[int(annotation['category_id'])]
         target = class_to_idx[id_name]
         date = image['date']
@@ -389,13 +425,15 @@ class DatasetMeta(data.Dataset):
             transform=None,
             train=False,
             aux_info=False,
-            dataset='inaturelist2021',
+            dataset='coco_generic',
             class_ratio=1.0,
             per_sample=1.0):
         self.aux_info = aux_info
         self.dataset = dataset
         if dataset in ['inaturelist2021','inaturelist2021_mini']:
             images, class_to_idx,images_info = find_images_and_targets(root,train,aux_info)
+        elif dataset in ['coco_generic']:
+            images, class_to_idx,images_info = find_images_and_targets(root,train,aux_info)
         elif dataset in ['inaturelist2017','inaturelist2018']:
             images, class_to_idx,images_info = find_images_and_targets_2017_2018(root,dataset,train,aux_info)
         elif dataset == 'cub-200':
@@ -427,7 +465,12 @@ class DatasetMeta(data.Dataset):
             path, target,aux_info = self.samples[index]
         else:
             path, target = self.samples[index]
-        img = open(path, 'rb').read() if self.load_bytes else Image.open(path).convert('RGB')
+
+        try:
+            img = open(path, 'rb').read() if self.load_bytes else Image.open(path).convert('RGB')
+        except:
+            img = Image.fromarray(np.zeros((224,224,3), dtype=np.uint8))
+
         if self.transform is not None:
             img = self.transform(img)
         if self.aux_info:

inference.py

@@ -7,6 +7,10 @@ from torch.autograd import Variable
 from torchvision.transforms import transforms
 import numpy as np
 import argparse
+from pycocotools.coco import COCO
+import requests
+import os
+from tqdm.auto import tqdm
 
 try:
     from apex import amp
@@ -34,24 +38,32 @@ def read_class_names(file_path):
     class_list = []
 
     for l in lines:
-        line = l.strip().split()
+        line = l.strip()
         # class_list.append(line[0])
-        class_list.append(line[1][4:])
+        class_list.append(line)
 
     classes = tuple(class_list)
     return classes
 
 
-class GenerateEmbedding:
-    def __init__(self, text_file):
-        self.text_file = text_file
+def read_class_names_coco(file_path):
+    dataset = COCO(file_path)
+    classes =  [dataset.cats[k]['name'] for k in sorted(dataset.cats.keys())]
+
+    with open("names.txt", 'w') as fp:
+        for c in classes:
+            fp.write(f"{c}\n")
 
+    return classes
+
+class GenerateEmbedding:
+    def __init__(self):
         self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
         self.model = AutoModel.from_pretrained("bert-base-uncased")
 
-    def generate(self):
+    def generate(self, text_file):
         text_list = []
-        with open(self.text_file, 'r') as f_text:
+        with open(text_file, 'r') as f_text:
             for line in f_text:
                 line = line.encode(encoding='UTF-8', errors='strict')
                 line = line.replace(b'\xef\xbf\xbd\xef\xbf\xbd', b' ')
@@ -69,57 +81,122 @@ class GenerateEmbedding:
 
 
 class Inference:
-    def __init__(self, config_path, model_path):
+    def __init__(self, config_path, model_path, names_path):
+
         self.config_path = config_path
         self.model_path = model_path
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        # self.classes = ("cat", "dog")
-        self.classes = read_class_names(r"D:\dataset\CUB_200_2011\CUB_200_2011\classes_custom.txt")
+        self.classes = read_class_names(names_path)
 
         self.config = model_config(self.config_path)
+
         self.model = build_model(self.config)
         self.checkpoint = torch.load(self.model_path, map_location='cpu')
-        self.model.load_state_dict(self.checkpoint['model'], strict=False)
+
+        if 'model' in self.checkpoint:
+            self.model.load_state_dict(self.checkpoint['model'], strict=False)
+        else:
+            self.model.load_state_dict(self.checkpoint, strict=False)
+            
         self.model.eval()
-        self.model.cuda()
+        self.model.to(self.device)
+        self.topk = 10
+        self.embedding_gen = GenerateEmbedding()
 
         self.transform_img = transforms.Compose([
-            transforms.Resize((224, 224), interpolation=Image.BILINEAR),
+            transforms.Resize((self.config.DATA.IMG_SIZE, self.config.DATA.IMG_SIZE), interpolation=Image.BILINEAR),
             transforms.ToTensor(), # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
             transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
         ])
 
-    def infer(self, img_path, meta_data_path):
-        _, _, meta = GenerateEmbedding(meta_data_path).generate()
-        meta = meta.cuda()
-        img = Image.open(img_path).convert('RGB')
+    def infer(self, img_path, meta_data_path, topk=None):
+
+        if isinstance(img_path, str):
+            if img_path.startswith("http"):
+                img = Image.open(requests.get(img_path, stream=True).raw).convert('RGB')
+            else:
+                img = Image.open(img_path).convert('RGB')
+        else:
+            img = img_path
+        
+        """
+        _, _, meta = self.embedding_gen(meta_data_path)
+        meta = meta.to(self.device)
+        """
+        meta = None
+
         img = self.transform_img(img)
         img.unsqueeze_(0)
-        img = img.cuda()
+        img = img.to(self.device)
         img = Variable(img).to(self.device)
         out = self.model(img, meta)
 
-        _, pred = torch.max(out.data, 1)
-        predict = self.classes[pred.data.item()]
-        # print(Fore.MAGENTA + f"The Prediction is: {predict}")
-        return predict
+        f = torch.nn.Softmax(dim=1)
+        y_pred = f(out)
+        indices = reversed(torch.argsort(y_pred, dim=1).squeeze().tolist())
+
+        if topk is not None:
+            predict = [{self.classes[idx] : y_pred.squeeze()[idx].cpu().item() for idx in indices[:topk]}]
+            return predict
+        else:
+            return {self.classes[idx] : y_pred.squeeze()[idx].cpu().item() for idx in indices}
 
 
 def parse_option():
     parser = argparse.ArgumentParser('MetaFG Inference script', add_help=False)
-    parser.add_argument('--cfg', type=str, default='D:/pycharmprojects/MetaFormer/configs/MetaFG_meta_bert_1_224.yaml', metavar="FILE", help='path to config file', )
+    parser.add_argument('--cfg', type=str, metavar="FILE", help='path to config file', default="configs/MetaFG_2_224.yaml")
     # easy config modification
-    parser.add_argument('--model-path', default='D:\pycharmprojects\MetaFormer\output\MetaFG_meta_1\cub_200\ckpt_epoch_92.pth', type=str, help="path to model data")
-    parser.add_argument('--img-path', default=r"D:\dataset\CUB_200_2011\CUB_200_2011\images\012.Yellow_headed_Blackbird\Yellow_Headed_Blackbird_0003_8337.jpg", type=str, help='path to image')
-    parser.add_argument('--meta-path', default=r"D:\dataset\CUB_200_2011\text_c10\012.Yellow_headed_Blackbird\Yellow_Headed_Blackbird_0003_8337.txt", type=str, help='path to meta data')
+    parser.add_argument('--model-path', type=str, help="path to model data", default="ckpt_epoch_12.pth")
+    parser.add_argument('--img-path', type=str, help='path to image')
+    parser.add_argument('--img-folder', type=str, help='path to image')
+    parser.add_argument('--meta-path', default="meta.txt", type=str, help='path to meta data')
+    parser.add_argument('--names-path', default="names_mf2.txt", type=str, help='path to meta data')
     args = parser.parse_args()
     return args
 
 
 if __name__ == '__main__':
     args = parse_option()
-    result = Inference(config_path=args.cfg,
-                       model_path=args.model_path).infer(img_path=args.img_path, meta_data_path=args.meta_path)
-    print("Predicted: ", result)
+    model = Inference(config_path=args.cfg,
+                       model_path=args.model_path,
+                       names_path=args.names_path)
+    
+    from glob import glob
+    glob_imgs = glob(os.path.join(args.img_folder, "*.jpg"))
+    out_dir = f"results_{os.path.splitext(os.path.basename(args.model_path))[0]}"
+    os.makedirs(out_dir, exist_ok=True)
+
+    for img in tqdm(glob_imgs):
+        try:
+            res = model.infer(img_path=img, meta_data_path=args.meta_path)
+        except KeyboardInterrupt:
+            break
+        except Exception as e:
+            print(e)
+            continue
+        
+        out = {}
+        out['preds'] = res
+
+        """
+        # Out is a list of (class, score). Return true/false if the top1 class is correct
+        out['top1_correct'] = '_'.join(res[0][1].split(' ')).lower() in os.path.basename(img).lower()
+
+        out['top5_correct'] = False
+        print(os.path.basename(img).lower())
+        for i in range(5):
+            out['top5_correct'] |= '_'.join(res[i][1].split(' ')).lower() in os.path.basename(img).lower()
+            print('_'.join(res[i][1].split(' ')).lower())
+            
+        out['top10_correct'] = False
+        for i in range(10):
+            out['top10_correct'] |= '_'.join(res[i][1].split(' ')).lower() in os.path.basename(img).lower()
+        """
+
+        # output json with inference results, use image basename 
+        # as filename
+        import json
+        with open(os.path.join(out_dir, os.path.splitext(os.path.basename(img))[0]+".json"), 'w') as fp:
+            json.dump(out, fp, indent=1)
 
 # Usage: python inference.py --cfg 'path/to/cfg' --model_path 'path/to/model' --img-path 'path/to/img' --meta-path 'path/to/meta'

lr_scheduler.py

@@ -21,7 +21,6 @@ def build_scheduler(config, optimizer, n_iter_per_epoch):
         lr_scheduler = CosineLRScheduler(
             optimizer,
             t_initial=num_steps,
-            t_mul=1.,
             lr_min=config.TRAIN.MIN_LR,
             warmup_lr_init=config.TRAIN.WARMUP_LR,
             warmup_t=warmup_steps,

main.py

@@ -2,7 +2,9 @@ import os
 import time
 import argparse
 import datetime
+import json
 import numpy as np
+from collections import defaultdict
 
 import torch
 import torch.backends.cudnn as cudnn
@@ -18,13 +20,23 @@ from lr_scheduler import build_scheduler
 from optimizer import build_optimizer
 from logger import create_logger
 from utils import load_checkpoint, save_checkpoint, get_grad_norm, auto_resume_helper, reduce_tensor,load_pretained
-from torch.utils.tensorboard import SummaryWriter
+
+have_wandb = False
+try:
+    import wandb
+    have_wandb = True
+except:
+    pass
+
+# TODO use torch.amp
 try:
     # noinspection PyUnresolvedReferences
     from apex import amp
 except ImportError:
     amp = None
 
+import logging
+logging.basicConfig(level=logging.INFO)
 
 def parse_option():
     parser = argparse.ArgumentParser('MetaFG training and evaluation script', add_help=False)
@@ -77,20 +89,19 @@ def parse_option():
                         help='dataset')
     parser.add_argument('--lr-scheduler-name', type=str,
                         help='lr scheduler name,cosin linear,step')
-    
+        
     parser.add_argument('--pretrain', type=str,
                         help='pretrain')
     
-    parser.add_argument('--tensorboard', action='store_true', help='using tensorboard')
-    
+    parser.add_argument('--wandb_job', type=str)
     
-    # distributed training
-    parser.add_argument("--local_rank", type=int, required=True, help='local rank for DistributedDataParallel')
-
     args, unparsed = parser.parse_known_args()
 
     config = get_config(args)
 
+    if have_wandb and int(config.LOCAL_RANK) == 0:
+        wandb.init(name = args.wandb_job, config=args)
+
     return args, config
 
 
@@ -98,14 +109,20 @@ def main(config):
     dataset_train, dataset_val, data_loader_train, data_loader_val, mixup_fn = build_loader(config)
     logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
     model = build_model(config)
+    if have_wandb and int(config.LOCAL_RANK) == 0:
+        wandb.config['model_config'] = config
     model.cuda()
     logger.info(str(model))
 
     optimizer = build_optimizer(config, model)
     if config.AMP_OPT_LEVEL != "O0":
         model, optimizer = amp.initialize(model, optimizer, opt_level=config.AMP_OPT_LEVEL)
-    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)
+    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[int(config.LOCAL_RANK)], broadcast_buffers=False)
     model_without_ddp = model.module
+
+    if have_wandb and int(config.LOCAL_RANK) == 0:
+        wandb.watch(model, log_freq=100)
+
     n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
     logger.info(f"number of params: {n_parameters}")
     if hasattr(model_without_ddp, 'flops'):
@@ -123,10 +140,15 @@ def main(config):
     max_accuracy = 0.0
     if config.MODEL.PRETRAINED:
         load_pretained(config,model_without_ddp,logger)
-        if config.EVAL_MODE:
-            acc1, acc5, loss = validate(config, data_loader_val, model)
-            logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
-            return
+
+        # Run initial validation
+        logger.info("Start validation (on init)")
+        acc1, acc5, loss, stats = validate(config, data_loader_val, model, limit=10)
+
+        with open(os.path.join(config.OUTPUT, f'val_init.json'), 'w') as fp:
+            json.dump(stats, fp, indent=1)
+
+        logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
 
     if config.TRAIN.AUTO_RESUME:
         resume_file = auto_resume_helper(config.OUTPUT)
@@ -143,11 +165,11 @@ def main(config):
     if config.MODEL.RESUME:
         logger.info(f"**********normal test***********")
         max_accuracy = load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, logger)
-        acc1, acc5, loss = validate(config, data_loader_val, model)
+        acc1, acc5, loss, stats = validate(config, data_loader_val, model)
         logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
         if config.DATA.ADD_META:
             logger.info(f"**********mask meta test***********")
-            acc1, acc5, loss = validate(config, data_loader_val, model,mask_meta=True)
+            acc1, acc5, loss, stats = validate(config, data_loader_val, model,mask_meta=True)
             logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
         if config.EVAL_MODE:
             return
@@ -165,18 +187,37 @@ def main(config):
             save_checkpoint(config, epoch, model_without_ddp, max_accuracy, optimizer, lr_scheduler, logger)
         
         logger.info(f"**********normal test***********")
-        acc1, acc5, loss = validate(config, data_loader_val, model)
+        acc1, acc5, loss, stats = validate(config, data_loader_val, model)
+
+        with open(os.path.join(config.OUTPUT, f'val_{epoch}.json'), 'w') as fp:
+            json.dump(stats, fp, indent=1)
+
         logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
         max_accuracy = max(max_accuracy, acc1)
         logger.info(f'Max accuracy: {max_accuracy:.2f}%')
         if config.DATA.ADD_META:
             logger.info(f"**********mask meta test***********")
-            acc1, acc5, loss = validate(config, data_loader_val, model,mask_meta=True)
+            acc1, acc5, loss, stats = validate(config, data_loader_val, model,mask_meta=True)
+
+            with open(os.path.join(config.OUTPUT, f'val_{epoch}_meta.json'), 'w') as fp:
+                json.dump(stats, fp, indent=1)
+
             logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
 #         data_loader_train.terminate()
+
+        if have_wandb and int(config.LOCAL_RANK) == 0:
+            wandb.run.summary["acc_top_1"] = acc1
+            wandb.run.summary["acc_top_5"] = acc5
+            wandb.run.summary["val_loss"] = loss
+
+            wandb.log({'val/acc1': acc1})
+            wandb.log({'val/acc5': acc5})
+            wandb.log({'val/loss': acc5})
+
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     logger.info('Training time {}'.format(total_time_str))
+
 def train_one_epoch_local_data(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler,tb_logger=None):
     model.train()
     if hasattr(model.module,'cur_epoch'):
@@ -261,6 +302,8 @@ def train_one_epoch_local_data(config, model, criterion, data_loader, optimizer,
             lr = optimizer.param_groups[0]['lr']
             memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
             etas = batch_time.avg * (num_steps - idx)
+            if have_wandb and int(config.LOCAL_RANK) == 0 and idx % 100 == 0:
+                wandb.log({"train/loss": loss_meter.val})
             logger.info(
                 f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t'
                 f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t'
@@ -271,7 +314,7 @@ def train_one_epoch_local_data(config, model, criterion, data_loader, optimizer,
     epoch_time = time.time() - start
     logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}")
 @torch.no_grad()
-def validate(config, data_loader, model, mask_meta=False):
+def validate(config, data_loader, model, mask_meta=False, limit=None):
     criterion = torch.nn.CrossEntropyLoss()
     model.eval()
 
@@ -280,8 +323,16 @@ def validate(config, data_loader, model, mask_meta=False):
     acc1_meter = AverageMeter()
     acc5_meter = AverageMeter()
 
+    stats = defaultdict(list)
+
     end = time.time()
+    
     for idx, data in enumerate(data_loader):
+        
+        if limit:
+            if idx > limit:
+                break
+
         if config.DATA.ADD_META:
             images,target,meta = data
             meta = [m.float() for m in meta]
@@ -314,6 +365,9 @@ def validate(config, data_loader, model, mask_meta=False):
         acc1_meter.update(acc1.item(), target.size(0))
         acc5_meter.update(acc5.item(), target.size(0))
 
+        for t in target:
+            stats[int(t.item())].append((acc1.item(), acc5.item(), loss.item()))
+
         # measure elapsed time
         batch_time.update(time.time() - end)
         end = time.time()
@@ -328,7 +382,8 @@ def validate(config, data_loader, model, mask_meta=False):
                 f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t'
                 f'Mem {memory_used:.0f}MB')
     logger.info(f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')
-    return acc1_meter.avg, acc5_meter.avg, loss_meter.avg
+
+    return acc1_meter.avg, acc5_meter.avg, loss_meter.avg, stats
 
 
 @torch.no_grad()
@@ -364,7 +419,7 @@ if __name__ == '__main__':
     else:
         rank = -1
         world_size = -1
-    torch.cuda.set_device(config.LOCAL_RANK)
+    torch.cuda.set_device(f'cuda:{config.LOCAL_RANK}')
     torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)
     torch.distributed.barrier()
 

models/MetaFG.py

@@ -54,7 +54,8 @@ class MetaFG(nn.Module):
                 qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,drop_path_rate=0.,
                 meta_dims=[],
                 only_last_cls=False,
-                use_checkpoint=False):
+                use_checkpoint=False,
+                **kwargs):
         super().__init__()
         self.only_last_cls = only_last_cls
         self.img_size = img_size