Add evaluation script

README.md

@@ -27,6 +27,7 @@ git clone https://github.com/hmichaeli/alias_free_convnets.git
 ```python
 from huggingface_hub import hf_hub_download
 import torch
+from torchvision import datasets, transforms
 from alias_free_convnets.models.convnext_afc import convnext_afc_tiny
 
 # baseline
@@ -51,6 +52,56 @@ afc_model = convnext_afc_tiny(
     )
 afc_model.load_state_dict(ckpt, strict=False)
 
+# evaluate model
+interpolation = transforms.InterpolationMode.BICUBIC
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+transform = transforms.Compose([
+    transforms.Resize(256, interpolation=interpolation),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD),
+])
+data_path = "/path/to/imagenet/val"
+dataset_val = datasets.ImageFolder(data_path, transform=transform)
+nb_classes = 1000
+sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+data_loader_val = torch.utils.data.DataLoader(
+            dataset_val, sampler=sampler_val,
+            batch_size=8,
+            num_workers=8,
+            drop_last=False
+        )
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+@torch.no_grad()
+def evaluate(data_loader, model, device):
+    model.eval()
+    correct = 0
+    total = 0
+    for batch_idx, (inputs, targets) in enumerate(data_loader):
+        inputs, targets = inputs.to(device), targets.to(device)
+        outputs = model(inputs)
+        _, predicted = outputs.max(1)
+        total += targets.size(0)
+        correct += predicted.eq(targets).sum().item()
+
+    acc = 100. * correct / total
+    print("Acc@1 {:.3f}".format(acc))
+
+
+
+print("evaluate baseline")
+base_model.to(device)
+test_stats = evaluate(data_loader_val, base_model, device)
+
+print("evaluate AFC")
+afc_model.to(device)
+test_stats = evaluate(data_loader_val, afc_model, device)
+
+
+
 ```