small change

app.py

@@ -1,13 +1,12 @@
 import torch
 import torch.nn as nn
-from torchvision import models
 import gradio as gr
 from PIL import Image
 import os
 from typing import List, Dict, Union
 
 from custom_transformer.vit import ViT
-from transforms import model_transforms
+from utils import model_transforms, get_pretrained_vit
 
 class GradioApp:
 
@@ -18,8 +17,7 @@ class GradioApp:
         custom = ViT().to(device).eval()
         custom.load_state_dict(torch.load('models/my_vit.pt', map_location=device))
         
-        pretrained = models.vit_b_16().to(device).eval()
-        pretrained.heads = nn.Linear(768, 3)
+        pretrained = get_pretrained_vit().to(device).eval()
         pretrained.load_state_dict(torch.load('models/pretrained_vit.pt', map_location=device))
 
         self.models: Dict[str, Union[str, nn.Module]] = {

transforms.py

@@ -1,12 +0,0 @@
-from torchvision import transforms, models
-from typing import Literal, Dict
-
-_weights = models.ViT_B_16_Weights.DEFAULT
-
-model_transforms: Dict[Literal['Custom', 'Pretrained'], transforms.Compose] = {
-    'Custom': transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor()
-    ]),
-    'Pretrained': _weights.transforms()
-}

utils.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+from torchvision import transforms, models
+from typing import Literal, Dict
+
+_weights = models.ViT_B_16_Weights.DEFAULT
+
+model_transforms: Dict[Literal['custom', 'pretrained'], Dict[Literal['train', 'val'], transforms.Compose]] = {
+    'custom': {
+        'train': transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.TrivialAugmentWide(),
+            transforms.ToTensor()
+        ]),
+        'val': transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor()
+        ])
+    },
+    'pretrained': {
+        'train': _weights.transforms(),
+        'val': _weights.transforms()
+    }
+}
+
+def get_pretrained_vit() -> models.VisionTransformer:
+    model = models.vit_b_16(weights='DEFAULT')
+    for parameter in model.parameters(): parameter.requires_grad = False
+    model.heads = nn.Linear(in_features=768, out_features=3)
+    return model