-
Notifications
You must be signed in to change notification settings - Fork 46
/
models.py
42 lines (35 loc) · 1.34 KB
/
models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import torch
import torch.nn as nn
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1, stride=1)
self.relu1 = nn.ReLU(inplace=True)
self.maxpool1 = nn.MaxPool2d(2)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=1)
self.relu2 = nn.ReLU(inplace=True)
self.maxpool2 = nn.MaxPool2d(2)
self.linear1 = nn.Linear(7*7*64, 200)
self.relu3 = nn.ReLU(inplace=True)
self.linear2 = nn.Linear(200, 10)
def forward(self, x):
out = self.maxpool1(self.relu1(self.conv1(x)))
out = self.maxpool2(self.relu2(self.conv2(out)))
out = out.view(out.size(0), -1)
out = self.relu3(self.linear1(out))
out = self.linear2(out)
return out
class SubstituteModel(nn.Module):
def __init__(self):
super(SubstituteModel, self).__init__()
self.linear1 = nn.Linear(28*28, 200)
self.relu1 = nn.ReLU(inplace=True)
self.linear2 = nn.Linear(200, 200)
self.relu2 = nn.ReLU(inplace=True)
self.linear3 = nn.Linear(200, 10)
def forward(self, x):
out = x.view(x.size(0), -1)
out = self.relu1(self.linear1(out))
out = self.relu2(self.linear2(out))
out = self.linear3(out)
return out