PrarthanaTS commited on
Commit
946a85c
·
1 Parent(s): 1f00ad2

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +70 -0
model.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.nn.functional as F
2
+ import torch.nn as nn
3
+
4
+ dropout = 0.01
5
+
6
+ class PrepBlock(nn.Module):
7
+ def __init__(self, dropout):
8
+ super(PrepBlock, self).__init__()
9
+ self.conv = nn.Sequential(
10
+ nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(3, 3), stride=1, padding=1, dilation=1, bias=False),
11
+ nn.ReLU(),
12
+ nn.BatchNorm2d(64),
13
+ nn.Dropout(dropout)
14
+ )
15
+
16
+ def forward(self, x):
17
+ return self.conv(x)
18
+
19
+ class ConvolutionBlock(nn.Module):
20
+ def __init__(self, in_channels, out_channels):
21
+ super(ConvolutionBlock, self).__init__()
22
+ self.conv = nn.Sequential(
23
+ nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(3, 3), stride=1, padding=1, bias=False),
24
+ nn.MaxPool2d(kernel_size=(2, 2)),
25
+ nn.BatchNorm2d(out_channels),
26
+ nn.ReLU()
27
+ )
28
+ def forward(self, x):
29
+ return self.conv(x)
30
+
31
+ class ResidualBlock(nn.Module):
32
+ def __init__(self, channels):
33
+ super(ResidualBlock, self).__init__()
34
+ self.residual = nn.Sequential(
35
+ nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=(3, 3), stride=1, padding=1, bias=False),
36
+ nn.BatchNorm2d(channels),
37
+ nn.ReLU(),
38
+ nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=(3, 3), stride=1, padding=1, bias=False),
39
+ nn.BatchNorm2d(channels),
40
+ nn.ReLU()
41
+ )
42
+ def forward(self, x):
43
+ return x + self.residual(x)
44
+
45
+
46
+ class Net(nn.Module):
47
+ def __init__(self):
48
+ super(Net, self).__init__()
49
+ self.prep = PrepBlock(dropout)
50
+ self.conv1 = ConvolutionBlock(64, 128)
51
+ self.R1 = ResidualBlock(128)
52
+ self.conv2 = ConvolutionBlock(128, 256)
53
+ self.conv3 = ConvolutionBlock(256, 512)
54
+ self.R2 = ResidualBlock(512)
55
+ self.maxpool = nn.MaxPool2d(kernel_size=(4, 4))
56
+ self.linear = nn.Linear(512, 10)
57
+
58
+ def forward(self, x):
59
+ x = self.prep(x)
60
+ x = self.conv1(x)
61
+ x = self.R1(x)
62
+ x = self.conv2(x)
63
+ x = self.conv3(x)
64
+ x = self.R2(x)
65
+ x = self.maxpool(x)
66
+ x = x.view(x.size(0), -1)
67
+ x = self.linear(x)
68
+ x = x.view(-1,10)
69
+ return F.log_softmax(x,dim=1)
70
+ return x