Skip to content

Instantly share code, notes, and snippets.

@giacaglia
Created December 8, 2019 00:55
Show Gist options
  • Save giacaglia/35e3d5a4a34c01a971ba0720b00a6f45 to your computer and use it in GitHub Desktop.
Save giacaglia/35e3d5a4a34c01a971ba0720b00a6f45 to your computer and use it in GitHub Desktop.
class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.fc = nn.Linear(7*7*32, num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.fc(out)
return out
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment