├── models ├── __init__.py ├── lenet.py ├── vgg.py ├── resnet.py └── googlenet.py ├── .gitignore ├── README.md ├── LICENSE └── main.py /models/__init__.py: -------------------------------------------------------------------------------- 1 | from .lenet import * 2 | from .resnet import * 3 | from .vgg import * 4 | from .googlenet import * -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # ignore data 2 | data/ 3 | 4 | # ignore cache 5 | models/__pycache__ 6 | 7 | # ignore saved model 8 | save -------------------------------------------------------------------------------- /models/lenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | class LeNet(nn.Module): 5 | def __init__(self): 6 | super(LeNet, self).__init__() 7 | 8 | self.conv1 = nn.Conv2d(3, 6, 5) 9 | self.conv2 = nn.Conv2d(6, 16, 5) 10 | 11 | self.relu = nn.ReLU() 12 | self.pool = nn.MaxPool2d(2,2) 13 | 14 | self.f1 = nn.Linear(16*5*5, 120) 15 | self.f2 = nn.Linear(120, 84) 16 | self.f3 = nn.Linear(84, 10) 17 | 18 | def forward(self, x): 19 | out = self.pool(self.relu(self.conv1(x))) 20 | out = self.pool(self.relu(self.conv2(out))) 21 | 22 | out = out.view(-1, 16*5*5) 23 | 24 | out = self.relu(self.f1(out)) 25 | out = self.relu(self.f2(out)) 26 | pred = self.relu(self.f3(out)) 27 | 28 | return pred 29 | 30 | 31 | 32 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pytorch-cifar10 2 | Classifying CIFAR10 dataset with popular DL computer vision models. 3 | 4 | # Requirements 5 | [Python 3x](https://www.python.org/) 6 | 7 | [PyTorch 1.0+](https://pytorch.org/get-started/locally/) 8 | 9 | CUDA and proper NVIDIA drivers (*optional, only if Nvidia GPU is available*) 10 | 11 | # Instructions 12 | `python main.py` 13 | 14 | `model` folder contains net architectures, just uncomment the preferred one in `main.py` 15 | 16 | # Current stats 17 | Model | GPU | Accuracy | Training Time 18 | --- | --- | --- | --- | 19 | [LeNet5](https://github.com/kanedaaaa/pytorch-cifar10/blob/main/models/lenet.py) | Tesla T4 | 67.15% | ... 20 | [ResNet](https://github.com/kanedaaaa/pytorch-cifar10/blob/main/models/resnet.py) | Tesla T4 | 76.14% | 21 min 21 | [VGG16](https://github.com/kanedaaaa/pytorch-cifar10/blob/main/models/vgg.py) | Tesla T4 | 78.60% | 6 min 22 | VGG19 | Tesla T4 | 78.51% | 7 min 23 | 24 | # Extras 25 | project is no longer maintained 26 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | This is free and unencumbered software released into the public domain. 2 | 3 | Anyone is free to copy, modify, publish, use, compile, sell, or 4 | distribute this software, either in source code form or as a compiled 5 | binary, for any purpose, commercial or non-commercial, and by any 6 | means. 7 | 8 | In jurisdictions that recognize copyright laws, the author or authors 9 | of this software dedicate any and all copyright interest in the 10 | software to the public domain. We make this dedication for the benefit 11 | of the public at large and to the detriment of our heirs and 12 | successors. We intend this dedication to be an overt act of 13 | relinquishment in perpetuity of all present and future rights to this 14 | software under copyright law. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 19 | IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 20 | OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 21 | ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | OTHER DEALINGS IN THE SOFTWARE. 23 | 24 | For more information, please refer to 25 | -------------------------------------------------------------------------------- /models/vgg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | vgg16 = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'] 6 | 7 | class VGG(nn.Module): 8 | def __init__(self, in_channels=3, num_classes=10): 9 | super(VGG, self).__init__() 10 | self.in_channels = in_channels 11 | self.conv = self._make_layer(vgg16) 12 | 13 | self.fc = nn.Sequential( 14 | nn.Linear(512*1*1, 4096), 15 | nn.ReLU(), 16 | nn.Dropout(p=0.5), 17 | nn.Linear(4096, 4096), 18 | nn.ReLU(), 19 | nn.Dropout(p=0.5), 20 | nn.Linear(4096, num_classes) 21 | ) 22 | 23 | def forward(self, x): 24 | x = self.conv(x) 25 | x = x.reshape(x.shape[0], -1) 26 | x = self.fc(x) 27 | return x 28 | 29 | def _make_layer(self, architecture): 30 | layers = [] 31 | in_channels = self.in_channels 32 | 33 | for x in architecture: 34 | if type(x) == int: 35 | out_channels = x 36 | 37 | layers += [ 38 | nn.Conv2d( 39 | in_channels=in_channels, 40 | out_channels=out_channels, 41 | kernel_size=(3, 3), 42 | stride=(1, 1), 43 | padding=(1, 1), 44 | ), 45 | nn.BatchNorm2d(x), 46 | nn.ReLU(), 47 | ] 48 | in_channels = x 49 | elif x == "M": 50 | layers += [nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))] 51 | 52 | return nn.Sequential(*layers) -------------------------------------------------------------------------------- /models/resnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | def conv3x3(in_channels, out_channels, stride=1): 6 | return nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) 7 | 8 | class ResidualBlock(nn.Module): 9 | expansion = 1 10 | 11 | def __init__(self, in_channels, out_channels, stride=1, downsample=None): 12 | super(ResidualBlock, self).__init__() 13 | self.conv1 = conv3x3(in_channels, out_channels, stride) 14 | self.bn1 = nn.BatchNorm2d(out_channels) 15 | self.relu = nn.ReLU(inplace=True) 16 | self.conv2 = conv3x3(out_channels, out_channels) 17 | self.bn2 = nn.BatchNorm2d(out_channels) 18 | self.downsample = downsample 19 | 20 | self.shortcut = nn.Sequential() 21 | if stride!=1 or in_channels != self.expansion * out_channels: 22 | self.shortcut = nn.Sequential( 23 | conv3x3(in_channels, self.expansion * out_channels, stride), 24 | nn.BatchNorm2d(self.expansion*out_channels)) 25 | 26 | 27 | def forward(self, x): 28 | out = self.relu(self.bn1(self.conv1(x))) 29 | out = self.bn2(self.conv2(out)) 30 | out += self.shortcut(x) 31 | out = self.relu(out) 32 | return out 33 | 34 | class ResNet(nn.Module): 35 | def __init__(self, block, layers, num_classes=10): 36 | super(ResNet, self).__init__() 37 | 38 | self.in_channels = 64 39 | self.conv = conv3x3(3, 64, 1) 40 | self.bn = nn.BatchNorm2d(64) 41 | self.relu = nn.ReLU(inplace=True) 42 | self.layer1 = self._make_layer(block, 64, layers[0], stride=1) 43 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 44 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 45 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 46 | 47 | self.fc = nn.Linear(512*block.expansion, num_classes) 48 | 49 | 50 | def _make_layer(self, block, out_channels, num_blocks, stride): 51 | strides = [stride] + [1]*(num_blocks-1) 52 | layers = [] 53 | for stride in strides: 54 | layers.append(block(self.in_channels, out_channels, stride)) 55 | self.in_channels = out_channels * block.expansion 56 | return nn.Sequential(*layers) 57 | 58 | def forward(self, x): 59 | out = self.relu(self.bn(self.conv(x))) 60 | out = self.layer1(out) 61 | out = self.layer2(out) 62 | out = self.layer3(out) 63 | out = self.layer4(out) 64 | out = F.avg_pool2d(out, 4) 65 | out = out.view(out.size(0), -1) 66 | out = self.fc(out) 67 | return out 68 | 69 | def resnet(): 70 | return ResNet(ResidualBlock, [2, 2, 2, 2]) 71 | 72 | 73 | -------------------------------------------------------------------------------- /models/googlenet.py: -------------------------------------------------------------------------------- 1 | class Inception(nn.Module): 2 | def __init__(self, in_planes, n1, n3_reduce, n3, n5_reduce, n5, pool_proj): 3 | super(Inception, self).__init__() 4 | 5 | self.block1 = nn.Sequential( 6 | nn.Conv2d(in_planes, n1, kernel_size=1), 7 | nn.BatchNorm2d(n1), 8 | nn.ReLU(True), 9 | ) 10 | 11 | self.block2 = nn.Sequential( 12 | nn.Conv2d(in_planes, n3_reduce, kernel_size=1), 13 | nn.BatchNorm2d(n3_reduce), 14 | nn.ReLU(True), 15 | 16 | nn.Conv2d(n3_reduce, n3, kernel_size=3, padding=1), 17 | nn.BatchNorm2d(n3), 18 | nn.ReLU(True), 19 | ) 20 | 21 | self.block3 = nn.Sequential( 22 | nn.Conv2d(in_planes, n5_reduce, kernel_size=1), 23 | nn.BatchNorm2d(n5_reduce), 24 | nn.ReLU(True), 25 | 26 | nn.Conv2d(n5_reduce, n5, kernel_size=3, padding=1), 27 | nn.BatchNorm2d(n5), 28 | nn.ReLU(True), 29 | 30 | nn.Conv2d(n5, n5, kernel_size=3, padding=1), 31 | nn.BatchNorm2d(n5), 32 | nn.ReLU(True), 33 | ) 34 | 35 | self.block4 = nn.Sequential( 36 | nn.MaxPool2d(3, stride=1, padding=1), 37 | nn.Conv2d(in_planes, pool_proj, kernel_size=1), 38 | nn.BatchNorm2d(pool_proj), 39 | nn.ReLU(True), 40 | ) 41 | 42 | def forward(self, x): 43 | y1 = self.block1(x) 44 | y2 = self.block2(x) 45 | y3 = self.block3(x) 46 | y4 = self.block4(x) 47 | 48 | return torch.cat([y1,y2,y3,y4], 1) 49 | 50 | 51 | 52 | class GoogLeNet(nn.Module): 53 | def __init__(self): 54 | super(GoogLeNet, self).__init__() 55 | 56 | self.pre_layers = nn.Sequential( 57 | nn.Conv2d(3, 192, kernel_size=3, padding=1), 58 | nn.BatchNorm2d(192), 59 | nn.ReLU(True), 60 | ) 61 | 62 | self.a3 = Inception(192, 64, 96, 128, 16, 32, 32) 63 | self.b3 = Inception(256, 128, 128, 192, 32, 96, 64) 64 | 65 | self.maxpool = nn.MaxPool2d(3, stride=2, padding=1) 66 | 67 | self.a4 = Inception(480, 192, 96, 208, 16, 48, 64) 68 | self.b4 = Inception(512, 160, 112, 224, 24, 64, 64) 69 | self.c4 = Inception(512, 128, 128, 256, 24, 64, 64) 70 | self.d4 = Inception(512, 112, 144, 288, 32, 64, 64) 71 | self.e4 = Inception(528, 256, 160, 320, 32, 128, 128) 72 | 73 | self.a5 = Inception(832, 256, 160, 320, 32, 128, 128) 74 | self.b5 = Inception(832, 384, 192, 384, 48, 128, 128) 75 | 76 | self.avg_pool = nn.AvgPool2d(8, stride=1) 77 | self.linear = nn.Linear(1024, 10) 78 | 79 | def forward(self, x): 80 | out = self.pre_layers(x) 81 | out = self.a3(out) 82 | out = self.b3(out) 83 | out = self.maxpool(out) 84 | out = self.a4(out) 85 | out = self.b4(out) 86 | out = self.c4(out) 87 | out = self.d4(out) 88 | out = self.e4(out) 89 | out = self.maxpool(out) 90 | out = self.a5(out) 91 | out = self.b5(out) 92 | out = self.avgpool(out) 93 | out = out.view(out.size(0), -1) 94 | out = self.linear(out) 95 | 96 | return out -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torchvision 4 | import torchvision.transforms as transforms 5 | import torch.optim as optim 6 | import matplotlib.pyplot as plt 7 | from tqdm.autonotebook import tqdm 8 | from models import * 9 | import numpy as np 10 | 11 | device = torch.device('cuda') 12 | print('Using device:', device) 13 | #print('GPU:', torch.cuda.get_device_name(0)) 14 | 15 | #net = resnet().to(device) 16 | #net = LeNet().to(device) 17 | #net = VGG().to(device) 18 | net = GoogLeNet().to(device) 19 | 20 | transform = transforms.Compose( 21 | [transforms.ToTensor(), 22 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] 23 | ) 24 | 25 | trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) 26 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=100, shuffle=True) 27 | 28 | testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) 29 | testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=True) 30 | 31 | classes = ('plane', 'car', 'bird', 'cat', 'deer', 32 | 'dog', 'frog', 'horse', 'ship', 'truck') 33 | 34 | 35 | criterion = nn.CrossEntropyLoss() 36 | optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) 37 | 38 | num_epochs = 7 39 | losses = [] 40 | batches = len(trainloader) 41 | 42 | def train(): 43 | print("TRAINING") 44 | for epoch in range(num_epochs): 45 | progress = tqdm(enumerate(trainloader), desc="Loss: ", total=batches) 46 | total_loss = 0 47 | for i, (inputs, labels) in progress: 48 | inputs, labels = inputs.to(device), labels.to(device) 49 | optimizer.zero_grad() 50 | output= net(inputs) 51 | loss = criterion(output, labels) 52 | 53 | loss.backward() 54 | optimizer.step() 55 | 56 | current_loss = loss.item() 57 | total_loss += current_loss 58 | progress.set_description("Loss: {:.4f}".format(total_loss/(i+1))) 59 | losses.append(total_loss/batches) 60 | print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/batches}") 61 | 62 | #torch.save(net, './save') 63 | 64 | 65 | def test(): 66 | total_correct = 0 67 | total_images = 0 68 | confusion_matrix = np.zeros([10,10], int) 69 | with torch.no_grad(): 70 | for inputs, labels in testloader: 71 | inputs, labels = inputs.to(device), labels.to(device) 72 | outputs = net(inputs) 73 | _, predicted = torch.max(outputs.data, 1) 74 | total_images += labels.size(0) 75 | total_correct += (predicted == labels).sum().item() 76 | for i, l in enumerate(labels): 77 | confusion_matrix[l.item(), predicted[i].item()] += 1 78 | 79 | model_accuracy = total_correct / total_images * 100 80 | print('Model accuracy on {0} test images: {1:.2f}%'.format(total_images, model_accuracy)) 81 | 82 | #net = torch.load("save") 83 | 84 | train() 85 | test() 86 | --------------------------------------------------------------------------------