[python] 파이썬 PyTorch에서 컨볼루션 신경망(CNN)을 구성하는 방법은?

22 Nov 2023

python

PyTorch는 딥러닝 모델을 구성하고 학습하는 동안 많은 지원 기능을 제공합니다. 이 중 하나는 CNN(컨볼루션 신경망)을 구성하는 데 사용할 수 있는 강력한 도구입니다.

다음은 PyTorch로 간단한 컨볼루션 신경망을 구성하는 코드 예시입니다.

import torch
import torch.nn as nn
import torch.optim as optim

# 컨볼루션 신경망 모델 정의
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(32 * 8 * 8, 10)  # 입력 크기에 맞게 수정

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.maxpool2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 모델 인스턴스 생성
model = CNN()

# 손실 함수와 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 학습과정
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 2000 == 1999:
            print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 2000}')
            running_loss = 0.0

print('학습이 완료되었습니다.')

이 코드에서는 nn.Conv2d() 함수를 사용하여 컨볼루션 레이어를 만들고, nn.MaxPool2d()를 사용하여 풀링 레이어를 만듭니다. 이후 nn.Linear()를 활용하여 완전 연결 레이어를 추가합니다. 이 모델은 이미지 분류와 같은 작업에 사용될 수 있습니다.

CNN 모델을 구성하는 방법에 대한 더 많은 내용은 PyTorch 공식 문서를 참고하십시오.

참고: PyTorch 공식 문서 - Convolutional Neural Network