-
Notifications
You must be signed in to change notification settings - Fork 2
/
resnet.py
114 lines (90 loc) · 3.99 KB
/
resnet.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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
# -*- coding: utf-8 -*-
"""
Created on Sat Feb 8 22:10:18 2020
---------------------------------------------------------------------
-- Author: Vigneashwara Pandiyan
---------------------------------------------------------------------
Utils file for visualization/ Plots
"""
#%%
import torch
import numpy as np
import torchvision.transforms as transforms
import torchvision
import torch.nn.functional as F
class baseBlock(torch.nn.Module):
expansion = 1
def __init__(self,input_planes,planes,stride=1,dim_change=None):
super(baseBlock,self).__init__()
#declare convolutional layers with batch norms
self.conv1 = torch.nn.Conv2d(input_planes,planes,stride=stride,kernel_size=3,padding=1)
self.bn1 = torch.nn.BatchNorm2d(planes)
self.conv2 = torch.nn.Conv2d(planes,planes,stride=1,kernel_size=3,padding=1)
self.bn2 = torch.nn.BatchNorm2d(planes)
self.dim_change = dim_change
def forward(self,x):
#Save the residue
res = x
output = F.relu(self.bn1(self.conv1(x)))
output = self.bn2(self.conv2(output))
if self.dim_change is not None:
res = self.dim_change(res)
output += res
output = F.relu(output)
return output
class bottleNeck(torch.nn.Module):
expansion = 4
def __init__(self,input_planes,planes,stride=1,dim_change=None):
super(bottleNeck,self).__init__()
self.conv1 = torch.nn.Conv2d(input_planes,planes,kernel_size=1,stride=1)
self.bn1 = torch.nn.BatchNorm2d(planes)
self.conv2 = torch.nn.Conv2d(planes,planes,kernel_size=3,stride=stride,padding=1)
self.bn2 = torch.nn.BatchNorm2d(planes)
self.conv3 = torch.nn.Conv2d(planes,planes*self.expansion,kernel_size=1)
self.bn3 = torch.nn.BatchNorm2d(planes*self.expansion)
self.dim_change = dim_change
def forward(self,x):
res = x
output = F.relu(self.bn1(self.conv1(x)))
output = F.relu(self.bn2(self.conv2(output)))
output = self.bn3(self.conv3(output))
if self.dim_change is not None:
res = self.dim_change(res)
output += res
output = F.relu(output)
return output
class ResNet(torch.nn.Module):
def __init__(self,block,num_layers,classes=4):
super(ResNet,self).__init__()
#according to research paper:
self.input_planes = 64
self.conv1 = torch.nn.Conv2d(3,64,kernel_size=3,stride=1,padding=1)
self.bn1 = torch.nn.BatchNorm2d(64)
self.layer1 = self._layer(block,64,num_layers[0],stride=1)
self.layer2 = self._layer(block,128,num_layers[1],stride=2)
self.layer3 = self._layer(block,256,num_layers[2],stride=2)
self.layer4 = self._layer(block,512,num_layers[3],stride=2)
self.averagePool = torch.nn.AvgPool2d(kernel_size=4,stride=1)
self.fc = torch.nn.Linear(512*block.expansion,classes)
def _layer(self,block,planes,num_layers,stride=1):
dim_change = None
if stride!=1 or planes != self.input_planes*block.expansion:
dim_change = torch.nn.Sequential(torch.nn.Conv2d(self.input_planes,planes*block.expansion,kernel_size=1,stride=stride),
torch.nn.BatchNorm2d(planes*block.expansion))
netLayers =[]
netLayers.append(block(self.input_planes,planes,stride=stride,dim_change=dim_change))
self.input_planes = planes * block.expansion
for i in range(1,num_layers):
netLayers.append(block(self.input_planes,planes))
self.input_planes = planes * block.expansion
return torch.nn.Sequential(*netLayers)
def forward(self,x):
x = F.relu(self.bn1(self.conv1(x)))
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = F.avg_pool2d(x,4)
x = x.view(x.size(0),-1)
x = self.fc(x)
return x