changed cnn architecture

bird_cnn/bird_cnn.py

@@ -27,14 +27,25 @@ class  SeparableConvolution(nn.Module):
 class SkipBlock(nn.Module):
     def __init__(self, c_in, c_out, kernel_size=3):
         super().__init__()
-        self.conv = SeparableConvolution(c_in=c_in, c_out=c_out, kernel_size=kernel_size)
+        self.conv = nn.Sequential(
+            nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2),
+            nn.BatchNorm2d(c_out),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
+            nn.BatchNorm2d(c_out),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
+            nn.BatchNorm2d(c_out),
+            nn.ReLU(inplace=True)
+        )
         self.conv_skip = nn.Sequential(
             nn.Conv2d(c_in, c_out, 1),
-            nn.BatchNorm2d(c_out)
+            nn.BatchNorm2d(c_out),
+            nn.ReLU(inplace=True)
         )
 
     def forward(self, x):
-        return(F.relu(self.conv_skip(x) + self.conv(x)))
+        return(F.relu(self.conv_skip(x) + self.conv(x), inplace=True))
 
 
 class Bird_CNN(nn.Module):
@@ -44,7 +55,7 @@ class Bird_CNN(nn.Module):
         self.model = nn.Sequential(
             nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
             nn.BatchNorm2d(c_hidden),
-            nn.ReLU(),
+            nn.ReLU(inplace=True),
 
             SkipBlock(c_in=c_hidden, c_out=c_hidden),
             SkipBlock(c_in=c_hidden, c_out=c_hidden),
@@ -56,7 +67,8 @@ class Bird_CNN(nn.Module):
             SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
             SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
 
-            SeparableConvolution(c_in=c_hidden*2, c_out=c_hidden*4, kernel_size=3),
+            nn.Conv2d(c_hidden*2, c_hidden*4, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
 
             nn.AdaptiveAvgPool2d((1, 1)),
             nn.Flatten(),

bird_cnn/server.py

@@ -14,7 +14,7 @@ from bird_cnn import Bird_CNN
 import threading
 
 BUILD_PATH = "./build_models"
-IMAGE_SIZE = 256
+IMAGE_SIZE = 64
 
 sem = threading.Semaphore(1) #adjust to performance
 
@@ -35,25 +35,13 @@ transform = transforms.Compose([
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model =Bird_CNN(c_in=3, c_hidden=64, c_out=7)
+model = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
 full_path = os.path.join(BUILD_PATH, "bird_cnn")
 model.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
 model.to(device)
 model.eval()
 
-app = FastAPI()
-
-origins = [
-    "*"
-]
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+app = FastAPI(root_path="/cnn-api")
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):

bird_cnn/train_bird.py

@@ -1,7 +1,9 @@
 import torch
 import torch.nn as nn
+import torchvision
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader, random_split
+import matplotlib.pyplot as plt
 
 from bird_cnn import Bird_CNN, sample, trainCNN
 
@@ -10,7 +12,7 @@ from enum import Enum
 from PIL import Image
 
 SAVE_PATH = "./saved_models"
-IMAGE_SIZE = 256
+IMAGE_SIZE = 64
 
 class bird_species(Enum):
     Common_Kingfisher = 0
@@ -34,13 +36,14 @@ val_size = len(dataset) - train_size
 
 train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
 
-train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
-val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
+
 
 device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
 print("Using device", device)
 
-model = Bird_CNN(c_in=3, c_hidden=64, c_out=7)
+model = Bird_CNN(c_in=3, c_hidden=16, c_out=7)
 model.to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
 loss_module = nn.CrossEntropyLoss()

bird_cnn/train_mnist.py

@@ -27,13 +27,13 @@ val_dataset = datasets.MNIST(
     transform=transform
 )
 
-train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
-val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)
+train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
 
 device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
 print("Using device", device)
 
-model = Bird_CNN(c_in=1, c_hidden=8, c_out=10)
+model = Bird_CNN(c_in=1, c_hidden=4, c_out=10)
 model.to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
 loss_module = nn.CrossEntropyLoss()