[Pytorch] 생초보의 파이토치 일기 - MNIST 손글씨 데이터 분류 99% 달성하기 2

지난 포스팅에서 CNN 모델을 이용해서 99%의 분류 성능을 달성하였다. 원하는 목표를 달성하였으니 좀 더 코드를 정리하고 내 나름대로 이쁘게 정리해보도록 하자. 일단 사진과 같이 코드 파일을 신경망을 구성하는 model.py와 학습을 하는 main.py을 나누자.

1. model.py

이 코드에서는 신경망을 정의하는 클래스들을 모아놓았다. 신경망은 지난 포스팅에서 정의한 모델 DNN, CNN 그대로 사용하였다. 먼저, DNN 모델 클래스이다.

# DNN 신경망 구성
class DNN(nn.Module) :
    def __init__(self):
        super(DNN, self).__init__()

        # Input shape = (?, 28, 28, 1) -> (?, 28 * 28) = (?, 784)
        # Dense 1    -> (?, 784) -> Weight[(784, 256)] -> (?, 256)
        self.fc1 = nn.Linear(784, 256)
        # Dense 2    -> (?, 256) -> Weight[(256, 128)] -> (?, 128)
        self.fc2 = nn.Linear(256, 128)
        # Dense 3    -> (?, 128) -> Weight[(128, 64)] -> (?, 64)
        self.fc3 = nn.Linear(128, 64)
        # Dense 4    -> (?, 64) -> Weight[(64, 10)] -> (?, 10) -> output shape
        self.fc4 = nn.Linear(64, 10)

    def forward(self, inputs):
        x = inputs.view(-1, 784)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)

        return F.log_softmax(x, dim=1)

그 다음으로 CNN 모델 클래스이다.

# CNN 신경망 구성
class CNN(nn.Module) :
    def __init__(self):
        super(CNN, self).__init__()

        # input shape = (?, 28, 28, 1)
        #   Conv     -> (?, 28, 28, 32)
        #   Pool     -> (?, 14, 14, 32)
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # input shape = (?, 14, 14, 32)
        #   Conv     -> (?, 14, 14, 64)
        #   Pool     -> (?, 7, 7, 64)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Dense 1    -> (?, 7*7*64) -> Weight[(7*7*64, 256)] -> (?, 256)
        self.fc1 = nn.Linear(7*7*64, 256)
        # Dense 2    -> (?, 256) -> Weight[(256, 128)] -> (?, 128)
        self.fc2 = nn.Linear(256, 128)
        # Dense 3    -> (?, 128) -> Weight[(128, 64)] -> (?, 64)
        self.fc3 = nn.Linear(128, 64)
        # Dense 4    -> (?, 64) -> Weight[(64, 10)] -> (?, 10) -> output shape
        self.fc4 = nn.Linear(64, 10)
    def forward(self, inputs):
        x = F.relu(self.conv1(inputs))
        x = self.maxpool1(x)
        x = F.relu(self.conv2(x))
        x = self.maxpool2(x)

        x = x.view(x.size(0), -1) # Flatten
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)

        return F.log_softmax(x, dim=1)

 

2. main.py

다음 코드 파일은 학습과 검증을 함께하는 실행 코드입니다. 먼저 argparse 라이브러리를 이용해서 인자를 파싱한다.

    # 인자 파싱
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, help='select model(DNN, CNN)')
    parser.add_argument('--batch_size', type=int, help='batch size')
    parser.add_argument('--epochs', type=int, help='epochs')
    parser.add_argument('--lr', type=float, help='learning rate')

    args = parser.parse_args()

파싱할 인자는 어떤 모델을 쓸 것인지, 배치 사이즈, 에폭 횟수, 학습률이다. 이렇게 인자를 미리 파싱해놓으면 터미널에서 아래의 명령어를 통해서 쉽게 실행을 할 수 있다.

python main.py --model CNN --batch_size 128 --epochs 20 --lr 0.01

위 명령어를 해석하면 모델은 CNN을 사용하고, 배치 사이즈는 128, 에폭은 20, 학습률은 0.01을 사용한다는 의미한다. 나중에 모델이나 다른 하이퍼파라미터를 명령어에 따라서 파싱할 때 아래와 같이 인스턴스를 하면 된다.

    # 하이퍼파라미터 정의
    batch_size = args.batch_size
    epochs = args.epochs
    lr = args.lr
    
    # 모델 인스턴스
    if args.model == "DNN" :
        model = DNN()
    if args.model == "CNN" :
        model = CNN()
    print("모델 정의 완료!")

이제 훈련 함수인 train과 평가 함수는 test를 정의한다.

def train(model, train_loader, device, optimizer, epoch, batch_size) :
    model.train()  # 모델을 train 모드로 변환
    train_loss = 0
    for batch_idx, (data, label) in enumerate(train_loader):
        data, label = data.to(device), label.to(device)
        optimizer.zero_grad()  # 그래디언트 초기화
        output = model(data)  # forward propagation
        loss = F.nll_loss(output, label)  # 손실 계산
        loss.backward()  # backward propagation
        optimizer.step()  # 가중치 업데이트
        if batch_idx % 50 == 0 :
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
        train_loss += loss.item()
    train_loss_list.append(train_loss / batch_size)

def test(model, device, test_loader) :
    model.eval() # 모델을 eval 모드로 변환
    test_loss = 0
    correct = 0
    with torch.no_grad() :
        for data, label in test_loader :
            data, label = data.to(device), label.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, label, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(label.view_as(pred)).sum().item()

        test_loss /= len(test_loader.dataset)
        test_loss_list.append(test_loss)

        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format
              (test_loss, correct, len(test_loader.dataset),
               100. * correct / len(test_loader.dataset)))

이제 아래의 반복문을 이용해서 학습을 진행하면 된다.

    # 최적화 알고리즘 정의
    optimizer = optim.SGD(params=model.parameters(), lr=lr)
    for epoch in range(epochs):
        train(model, train_loader, device, optimizer, epoch, batch_size)
        test(model, device, test_loader)
    plot_loss_acc(train_loss_list, test_loss_list, epochs)