Ramendra commited on
Commit
bee69fc
1 Parent(s): 9f5f4c1

Upload 3 files

Browse files
Files changed (3) hide show
  1. cat_dog_classifier.pth +3 -0
  2. model.py +28 -0
  3. requirements.txt +3 -0
cat_dog_classifier.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:57f43885cf2dac0be082e7107b921821cc02641999d75e7c96068cc918c12c0d
3
+ size 228053947
model.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from torch import manual_seed, nn
3
+ from torchvision import transforms, models
4
+
5
+ def create_model_alexnet(num_classes:int=2, seed:int=42):
6
+ """Creates model and transforms.
7
+
8
+ Args:
9
+ num_classes (int, optional): number of classes in the classifier head, defaults to 2.
10
+ seed (int, optional): random seed value. Defaults to 42.
11
+
12
+ Returns:
13
+ model (torch.nn.Module): Alexnet model.
14
+ transforms (torchvision.transforms): Alexnet image transforms.
15
+ """
16
+ # Create Alexnet pretrained weights, transforms and model
17
+ weights = models.AlexNet_Weights.IMAGENET1K_V1.DEFAULT
18
+ auto_transform = weights.transforms()
19
+ model_alexnet = models.alexnet(weights=weights)
20
+
21
+ # Freeze all layers in base model
22
+ for param in model_alexnet.parameters():
23
+ param.requires_grad = False
24
+
25
+ # Change classifier head with random seed for reproducibility
26
+ manual_seed(seed)
27
+ model_alexnet.classifier[6] = nn.Linear(4096, out_features=num_classes)
28
+ return model_alexnet, auto_transform
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch==2.0.1+cu118
2
+ torchvision==0.15.2+cu118
3
+ gradio==3.47.1