이제 데이터셋을 학습시킬 모델의 네트워크를 설계하기 위한 클래스를 생성합니다.
코드 6-10 모델의 네트워크 클래스
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.cnn1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=0) ------ 2D 합성곱층이 적용됩니다. 이때 입력 형태는 (3, 224, 224)가 되며 출력 형태는 (weight-kernel_size+1)/stride에 따라 (16, 220, 220)이 됩니다.
self.relu1 = nn.ReLU() ------ ReLU 활성화 함수입니다.
self.maxpool1 = nn.MaxPool2d(kernel_size=2) ------ 최대 풀링이 적용됩니다. 적용 이후 출력 형태는 220/2가 되어 (16, 110, 110)입니다.
self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=0) ------ 또다시 2D 합성곱층이 적용되며 출력 형태는 (32, 106, 106)입니다.
self.relu2 = nn.ReLU()
self.maxpool2 = nn.MaxPool2d(kernel_size=2) ------ 최대 풀링이 적용되며 출력 형태는 (32, 53, 53)입니다.
self.fc1 = nn.Linear(32*53*53, 512)
self.relu5 = nn.ReLU()
self.fc2 = nn.Linear(512, 2)
self.output = nn.Softmax(dim=1)
def forward(self, x):
out = self.cnn1(x)
out = self.relu1(out)
out = self.maxpool1(out)
out = self.cnn2(out)
out = self.relu2(out)
out = self.maxpool2(out)
out = out.view(out.size(0), -1) ------ 완전연결층에 데이터를 전달하기 위해 데이터 형태를 1차원으로 바꿉니다.
out = self.fc1(out)
out = self.fc2(out)
out = self.output(out)
return out