이렇게 작성해서 돌리고 있는데 혹시 틀린 문법이 있을까요?ㅠㅠ



import torch

import torch.nn as nn


class Block(nn.Module):

    def __init__(self, c_in, c_out, stride=1):

        super(Block, self).__init__()


        self.c_in = c_in

        self.c_mid = c_out // 4

        self.c_out = c_out


        if self.c_in != self.c_out:

            self.conv_1 = nn.Conv2d(self.c_in, self.c_mid, 1, stride=stride)

        else:

            self.conv_1 = nn.Conv2d(self.c_in, self.c_mid, 1, 1)


        self.conv_2 = nn.Conv2d(self.c_mid, self.c_mid, 3, 1, 1)

        self.conv_3 = nn.Conv2d(self.c_mid, self.c_out, 1, 1)

        self.conv_map = nn.Conv2d(self.c_in, self.c_out, 1, stride=stride)

        self.bn1 = nn.BatchNorm2d(self.c_mid)

        self.bn2 = nn.BatchNorm2d(self.c_mid)

        self.bn3 = nn.BatchNorm2d(self.c_out)

        self.relu = nn.ReLU()



    def forward(self, x):

        if self.c_in != self.c_out:

            mapping = self.bn3(self.conv_map(x))

        else:

            mapping = x

        x = self.relu(self.bn1(self.conv_1(x)))

        x = self.relu(self.bn2(self.conv_2(x)))

        x = self.bn3(self.conv_3(x))

        x += mapping

        x = self.relu(x)


        return x



class ResNet50(nn.Module):

    def __init__(self, num_class):

        super(ResNet50, self).__init__()

                                            # size: (3,128,128)

        self.layer1 = nn.Sequential(

            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),

            nn.BatchNorm2d(64),             # size: (64,64,64)

            nn.ReLU()

        )

        self.max = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)   # size: (64,32,32)

        self.layer2 = nn.Sequential(

            Block(64, 256, 1),  # size: (256,32,32)

            Block(256, 256),          # size: (256,32,32)

            Block(256, 256)           # size: (256,32,32)

        )


        self.layer3 = nn.Sequential(            # size: (512,16,16)

            Block(256, 512, 2), # size: (512,16,16)

            Block(512, 512),          # size: (512,16,16)

            Block(512, 512),          # size: (512,16,16)

            Block(512, 512)           # size: (512,16,16)

        )


        self.layer4 = nn.Sequential(

            Block(512, 1024, 2),    # size: (1024,8,8)

            Block(1024, 1024),

            Block(1024, 1024),

            Block(1024, 1024),

            Block(1024, 1024),

            Block(1024, 1024)             # size: (1024,8,8)


        )


        self.layer5 = nn.Sequential(

            Block(1024, 2048, 2),   # size: (2048,4,4)

            Block(2048, 2048),

            Block(2048, 2048)             # size: (2048,8,8)

        )


        self.layer6 = nn.Sequential(

            nn.AdaptiveAvgPool2d((1, 1)),

            nn.Flatten(),

            nn.Linear(2048, num_class)

        )


    def forward(self, x):


        x = self.layer1(x)

        x = self.max(x)

        x = self.layer2(x)

        x = self.layer3(x)

        x = self.layer4(x)

        x = self.layer5(x)

        x = self.layer6(x)


        return x