Upload main.py

main.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""main.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/#fileId=https%3A//huggingface.co/spaces/neelimapreeti297/panda_cat_dog_classification/blob/main/main.ipynb
+"""
+
+#import libraries
+import torch
+from torchvision import datasets, transforms
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+
+#define the data transforms
+
+transform = transforms.Compose([
+  transforms.Resize((224,224)),
+  transforms.ToTensor(),
+  transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
+  ])
+
+#insert the datasets
+
+train_dataset = ImageFolder('./data/train', transform=transform)
+test_dataset =ImageFolder('./data/test', transform=transform)
+
+# make cnn model
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.fc1 = nn.Linear(16 * 53 * 53, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 3)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.pool(x)
+        x = self.conv2(x)
+        x = self.pool(x)
+        x = x.view(-1, 16 * 53 * 53)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+        return x
+
+batch_size = 8
+
+train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
+
+model = CNN()
+loss_function = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+#Train the model
+
+for epoch in range(10):
+    for i, (images, labels) in enumerate(train_loader):
+
+        outputs = model(images)
+
+        loss = loss_function(outputs, labels)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if i % 200 == 0:
+            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, 10, i + 1, len(train_loader), loss.item()))
+
+#iterate over the test data
+
+correct = 0
+total = 0
+for i, (images, labels) in enumerate(test_loader):
+  outputs = model(images)
+
+  _, predicted = torch.max(outputs.data, 1)
+  correct += (predicted == labels).sum().item()
+  total += labels.size(0)
+
+#calculate the accuracy
+accuracy = 100 * correct / total
+print('Accuracy: {}%' .format(accuracy))
+
+model_scripted = torch.jit.script(model)
+model_scripted.save('./models/cat_dog_cnn.pt')
+