Skip to content
Snippets Groups Projects
Commit b488badd authored by Mina Moshfegh's avatar Mina Moshfegh
Browse files

Upload New File

parent 38f1e902
No related branches found
No related tags found
No related merge requests found
import torch
import torch.nn as nn
import torch.nn.functional as F
# This is a basic CNN architecture suitable for smaller images,
# especially used for MNIST in many adversarial training setups.
class SmallCNN(nn.Module):
def __init__(self, num_channels=1, num_classes=10):
super().__init__()
# Four convolutional layers for feature extraction.
# Typically used for grayscale MNIST, so default num_channels=1.
self.conv1 = nn.Conv2d(num_channels, 32, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.conv4 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
# Then three fully-connected layers for classification.
self.fc1 = nn.Linear(64 * 7 * 7, 200)
self.fc2 = nn.Linear(200, 200)
self.fc3 = nn.Linear(200, num_classes)
def forward(self, x):
# Pass through two conv layers, each with relu,
# then do a 2x2 max pool.
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
# Another pair of conv+relu, then max pool.
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
x = F.max_pool2d(x, 2)
# Flatten the feature maps into a 1D vector.
x = x.view(x.size(0), -1)
# Pass through FC layers with relu, except final layer.
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
# Apply softmax for classification outputs.
x = F.softmax(x, dim=1)
return x
def forward_features(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
x = F.max_pool2d(x, 2)
x = x.view(x.size(0), -1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
return x
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment