drawing/model.py at main · Ims2oha/drawing · 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
44
45
import torch
from torch import nn

from data import NUM_CLASSES

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.Conv_block1 = nn.Sequential(nn.Conv2d(1,16,3,padding=1),
                                        nn.BatchNorm2d(16),
                                        nn.ReLU(),
                                        nn.Conv2d(16,16,3,padding=1),
                                        nn.BatchNorm2d(16),
                                        nn.ReLU())
        self.Maxpool1 = nn.MaxPool2d(2)
        self.Conv_1x1 = nn.Sequential(nn.Conv2d(16,64,1),
                                      nn.BatchNorm2d(64),
                                      nn.ReLU(),
                                      nn.Conv2d(64,64,1))
        self.Conv_block2 = nn.Sequential(nn.Conv2d(64,64,3,padding=1),
                                        nn.BatchNorm2d(64),
                                        nn.ReLU(),
                                        nn.Conv2d(64,64,3,padding=1),
                                        nn.BatchNorm2d(64),
                                        nn.ReLU(),
                                        nn.Conv2d(64,64,3,padding=1),
                                        nn.BatchNorm2d(64),
                                        nn.ReLU())
        self.Maxpool2 = nn.MaxPool2d(2)#7x7
        self.Conv_1x1_2 = nn.Sequential(nn.Conv2d(64,128,1),
                                        nn.BatchNorm2d(128),
                                        nn.ReLU())#kernel_size 1인 CNN 즉 1x1CNN을 쓸려면 1x1으로 채널 축소 -> 3x3 CNN 통과 1x1으로 채널 증가하는 방식으로 증가 하고 축소해도 괜찮음
        self.fc1 = nn.Sequential(nn.Linear(128*7*7,1024),
                                 nn.ReLU(),
                                 nn.Linear(1024,NUM_CLASSES))
    def forward(self,x):
        x = self.Conv_block1(x)
        x = self.Maxpool1(x)
        x = self.Conv_1x1(x)
        x = self.Conv_block2(x)
        x = self.Maxpool2(x)
        x = self.Conv_1x1_2(x)
        x = torch.flatten(x,start_dim=1)
        x = self.fc1(x)
        return x