Skip to content

Instantly share code, notes, and snippets.

@AyoubOuddah
Last active November 26, 2019 17:30
Show Gist options
  • Save AyoubOuddah/62c3762340f6f226bf9fa1a34ee5256f to your computer and use it in GitHub Desktop.
Save AyoubOuddah/62c3762340f6f226bf9fa1a34ee5256f to your computer and use it in GitHub Desktop.
CNN on Fashion-MNIST dataset with Pytorch
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.optim
import torch.utils.data
CUDA = False
#----------------------------- NET CLASS -------------------------------#
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
#defining the layers
self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=5) # 28 X 28 --> 24 X 24 --> 12 X 12
self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=5) #12 X 12 --> 8 X 8 --> 4 X 4
self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=64) # 4 X 4 X 32 --> 64
self.fc2 = nn.Linear(in_features=64, out_features=10) # 64 --> 10
#forward pass
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=2, stride=2)
x = F.max_pool2d(F.relu(self.conv2(x)), kernel_size=2, stride=2)
x = x.view(-1, 16 * 4 * 4 )
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
def main(params):
#get the train data download it if needed
train = datasets.FashionMNIST(
root = './data/FashionMNIST',
train = True,
download = True,
transform = transforms.Compose([ transforms.ToTensor()]))
#get the test data download it if needed
test = datasets.FashionMNIST(
root = './data/FashionMNIST',
train = False,
download = True,
transform = transforms.Compose([ transforms.ToTensor()]))
#devide the train data into batchs and shaffle them
trainset = torch.utils.data.DataLoader(train, batch_size = params.batch_size,
shuffle = True, num_workers = 2)
#devide the train data into batchs and shaffle them
testset = torch.utils.data.DataLoader(test, batch_size = params.batch_size,
shuffle = True, num_workers = 2)
net = Net() #creating the net
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), params.lr) #creating the optimizer
if CUDA: #IF CUDA is enabled transforme the model into the GPU
net = net.cuda()
criterion = criterion.cuda()
#train the model
for i in range(params.epochs):
print("=================\n=== EPOCH "+str(i+1)+" =====\n=================\n")
for data in trainset:
X, y = data
if CUDA:
X = X.cuda()
y = y.cuda()
net.zero_grad()
output = net(X)
loss = criterion(output, y)
loss.backward()
optimizer.step()
print(loss)
#test the model
correct = 0
total = 0
with torch.no_grad():
for data in testset:
X, y = data
if CUDA:
X = X.cuda()
y = y.cuda()
output = net(X)
#print(output)
for idx, i in enumerate(output):
#print(torch.argmax(i), y[idx])
if torch.argmax(i) == y[idx]:
correct += 1
total += 1
print("Accuracy: ", round(correct/total, 3))
if __name__ == '__main__':
# Paramètres en ligne de commande
parser = argparse.ArgumentParser()
#parser.add_argument('--path', default='/tmp/datasets/mnist', type=str, metavar='DIR', help='path to dataset')
parser.add_argument('--epochs', default=5, type=int, metavar='N', help='number of total epochs to run')
parser.add_argument('--batch-size', default=128, type=int, metavar='N', help='mini-batch size (default: 128)')
parser.add_argument('--lr', default=0.1, type=float, metavar='LR', help='learning rate')
parser.add_argument('--cuda', dest='cuda', action='store_true', help='activate GPU acceleration')
args = parser.parse_args()
if args.cuda:
CUDA = True
#cudnn.benchmark = True
main(args)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment