PyTorch-practice/ResNet.py at master · IshantDere/PyTorch-practice · GitHub

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
# imports
import torch
from torch import nn

# created convolution block
class Conv(nn.Module):
    def __init__(self, in_channels,
                 out_channels,
                 kernal_size = (3,3),
                 stride = (1, 1),
                 padding = 1):
        super().__init__()

        self.conv = Conv(in_channels,
                         out_channels,
                         kernal_size,
                         stride,
                         padding)
        self.norm = nn.BatchNorm2d(in_channels)

        self.act = nn.ReLU()

    def forward(self, x):
        x = self.norm(x)
        x = self.conv(x)
        x = self.act(x)
        return x

    # created a ResNet architecture

    class ResnetBlock(nn.Module):
    def __init__(self, in_channels,
                 out_channels):
        super().__init__()

        self.conv = Conv(in_channels,
                         out_channels)

    def forward(self, x):
        residual = x
        x = self.conv(x)
        x = x + residual
        return x