Skip to content

Instantly share code, notes, and snippets.

@mohcinemadkour
Created July 25, 2020 05:44
Show Gist options
  • Save mohcinemadkour/f272c8ec380d1712ad456d03a30a8ae0 to your computer and use it in GitHub Desktop.
Save mohcinemadkour/f272c8ec380d1712ad456d03a30a8ae0 to your computer and use it in GitHub Desktop.
class Flatten(nn.Module):
def __init__(self):
super(Flatten, self).__init__()
def forward(self, x):
x = x.view(x.size(0), -1)
return x
class LeNet(nn.Module):
def __init__(self, droprate=0.5):
super(LeNet, self).__init__()
self.model = nn.Sequential()
self.model.add_module('conv1', nn.Conv2d(1, 20, kernel_size=5, padding=2))
self.model.add_module('dropout1', nn.Dropout2d(p=droprate))
self.model.add_module('maxpool1', nn.MaxPool2d(2, stride=2))
self.model.add_module('conv2', nn.Conv2d(20, 50, kernel_size=5, padding=2))
self.model.add_module('dropout2', nn.Dropout2d(p=droprate))
self.model.add_module('maxpool2', nn.MaxPool2d(2, stride=2))
self.model.add_module('flatten', Flatten())
self.model.add_module('dense3', nn.Linear(50*7*7, 500))
self.model.add_module('relu3', nn.ReLU())
self.model.add_module('dropout3', nn.Dropout(p=droprate))
self.model.add_module('final', nn.Linear(500, 10))
def forward(self, x):
return self.model(x)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment