Spaces:
Sleeping
Sleeping
Create custom_resnet.py
Browse files- custom_resnet.py +104 -0
custom_resnet.py
ADDED
@@ -0,0 +1,104 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch.nn as nn
|
2 |
+
import torch.nn.functional as F
|
3 |
+
|
4 |
+
|
5 |
+
class BasicBlock(nn.Module):
|
6 |
+
|
7 |
+
def __init__(self, in_planes, planes, stride=1):
|
8 |
+
super(BasicBlock, self).__init__()
|
9 |
+
self.conv1 = nn.Conv2d(
|
10 |
+
in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
|
11 |
+
)
|
12 |
+
self.bn1 = nn.BatchNorm2d(planes)
|
13 |
+
self.conv2 = nn.Conv2d(
|
14 |
+
planes, planes, kernel_size=3, stride=1, padding=1, bias=False
|
15 |
+
)
|
16 |
+
self.bn2 = nn.BatchNorm2d(planes)
|
17 |
+
|
18 |
+
self.shortcut = nn.Sequential()
|
19 |
+
|
20 |
+
def forward(self, x):
|
21 |
+
out = F.relu(self.bn1(self.conv1(x)))
|
22 |
+
out = self.bn2(self.conv2(out))
|
23 |
+
out += self.shortcut(x)
|
24 |
+
out = F.relu(out)
|
25 |
+
return out
|
26 |
+
|
27 |
+
|
28 |
+
class CustomBlock(nn.Module):
|
29 |
+
def __init__(self, in_channels, out_channels):
|
30 |
+
super(CustomBlock, self).__init__()
|
31 |
+
|
32 |
+
self.inner_layer = nn.Sequential(
|
33 |
+
nn.Conv2d(
|
34 |
+
in_channels=in_channels,
|
35 |
+
out_channels=out_channels,
|
36 |
+
kernel_size=3,
|
37 |
+
stride=1,
|
38 |
+
padding=1,
|
39 |
+
bias=False,
|
40 |
+
),
|
41 |
+
nn.MaxPool2d(kernel_size=2),
|
42 |
+
nn.BatchNorm2d(out_channels),
|
43 |
+
nn.ReLU(),
|
44 |
+
)
|
45 |
+
|
46 |
+
self.res_block = BasicBlock(out_channels, out_channels)
|
47 |
+
|
48 |
+
def forward(self, x):
|
49 |
+
x = self.inner_layer(x)
|
50 |
+
r = self.res_block(x)
|
51 |
+
|
52 |
+
out = x + r
|
53 |
+
|
54 |
+
return out
|
55 |
+
|
56 |
+
|
57 |
+
class CustomResNet(nn.Module):
|
58 |
+
def __init__(self, num_classes=10):
|
59 |
+
super(CustomResNet, self).__init__()
|
60 |
+
|
61 |
+
self.prep_layer = nn.Sequential(
|
62 |
+
nn.Conv2d(
|
63 |
+
in_channels=3,
|
64 |
+
out_channels=64,
|
65 |
+
kernel_size=3,
|
66 |
+
stride=1,
|
67 |
+
padding=1,
|
68 |
+
bias=False,
|
69 |
+
),
|
70 |
+
nn.BatchNorm2d(64),
|
71 |
+
nn.ReLU(),
|
72 |
+
)
|
73 |
+
|
74 |
+
self.layer_1 = CustomBlock(in_channels=64, out_channels=128)
|
75 |
+
|
76 |
+
self.layer_2 = nn.Sequential(
|
77 |
+
nn.Conv2d(
|
78 |
+
in_channels=128,
|
79 |
+
out_channels=256,
|
80 |
+
kernel_size=3,
|
81 |
+
stride=1,
|
82 |
+
padding=1,
|
83 |
+
bias=False,
|
84 |
+
),
|
85 |
+
nn.MaxPool2d(kernel_size=2),
|
86 |
+
nn.BatchNorm2d(256),
|
87 |
+
nn.ReLU(),
|
88 |
+
)
|
89 |
+
|
90 |
+
self.layer_3 = CustomBlock(in_channels=256, out_channels=512)
|
91 |
+
|
92 |
+
self.max_pool = nn.Sequential(nn.MaxPool2d(kernel_size=4))
|
93 |
+
|
94 |
+
self.fc = nn.Linear(512, num_classes)
|
95 |
+
|
96 |
+
def forward(self, x):
|
97 |
+
x = self.prep_layer(x)
|
98 |
+
x = self.layer_1(x)
|
99 |
+
x = self.layer_2(x)
|
100 |
+
x = self.layer_3(x)
|
101 |
+
x = self.max_pool(x)
|
102 |
+
x = x.view(x.size(0), -1)
|
103 |
+
x = self.fc(x)
|
104 |
+
return F.log_softmax(x,dim=1)
|