[Pytorch] Pytorch 기초 사용법 정리

double-oh·2021년 1월 28일
0

Pytorch 프레임워크 기초적인 사용법을 정리해보았습니다.

1. cpu, gpu에 텐서 할당

  • pytorch는 tensorflow와 다르게 gpu에 자동으로 할당해주지 않고, gpu 할당을 선언해주어야한다.
import torch

if __name__ == '__main__':

    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    x = torch.tensor([1., 2.]).to(device)
  • 2개이상 gpu 병렬연산 예시
model = Model(data)
if torch.cuda.device_count() > 1:
  model = nn.DataParallel(model)

model.to(device)

2. torch를 통해 데이터 로딩

  • torch의 dataloader를 이용하여 데이터 로딩 및 transform을 이용하여 전처리를 수행할 수 있다.
import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt


if __name__ == '__main__':

    # MNIST train 데이터 가져오기
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data/', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor()])),
        batch_size=32
    )

    # MNIST test 데이터 가져오기
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data/', train=False, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor()])),
        batch_size=32
    )

    # 학습 이미지 출력
    train_image, train_label = next(iter(train_loader))
    plt.imshow(train_image[0][0], 'gray')
    plt.show()

    test_image, test_label = next(iter(train_loader))
    plt.imshow(test_image[0][0], 'gray')
    plt.show()

3. Convolution 레이어로 MNIST 이미지 분류

  • Conv2d 레이어 2층과 FullyConnected 레이어 2층을 쌓은 모델을 사용한다.
  • loss함수는 CrossEntropy, optimizer는 Adam을 이용한다.
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms


class CnnNet(nn.Module):
    def __init__(self):
        super(CnnNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=20, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(in_channels=20, out_channels=50, kernel_size=5, stride=1)
        self.fc1 = nn.Linear(800, 100)  # 800 = ((((28 - 5 + 1) / 2) - 5 + 1) / 2) * ((((28 - 5 + 1) / 2) - 5 + 1) / 2) * 50
        self.fc2 = nn.Linear(100, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)

        x = x.view(-1, 800)
        x = F.relu(self.fc1(x))
        x = F.log_softmax(self.fc2(x), dim=1)

        return x


if __name__ == '__main__':
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # input 채널은 1, output 채널은 20, kernel 크기는 (3,3), stride는 (1,1)
    layer = nn.Conv2d(in_channels=1, out_channels=20, kernel_size=3, stride=1).to(device)

    print(layer)

    # 초기 weight 값
    print(type(layer.weight))

    # weight를 numpy로 만들려면 detach가 필요
    print(layer.weight.detach().cpu().numpy().shape)

    # MNIST train 데이터 가져오기
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data/', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor()])),
        batch_size=32
    )

    # 모델 정의
    model = CnnNet().to(device)

    # 학습 옵티마이저
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # loss 함수 정의
    criterion = nn.CrossEntropyLoss()

    # 학습
    EPOCHS = 10
    for epoch in range(EPOCHS):
        model.train()

        total_loss = 0

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            # loss 총합 계산
            total_loss += loss.item()

        print('epoch: {} loss: {:.4f}'.format(epoch + 1, total_loss))

    # MNIST test 데이터 가져오기
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data/', train=False, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor()])),
        batch_size=32
    )

    # 예측
    model.eval()

    test_loss = 0
    correct = 0

    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)

            output = model(data)

            test_loss += criterion(output, target).item()

            pred = output.argmax(dim=1, keepdim=True)

            correct += pred.eq(target.view_as(pred)).sum().item()

    print('loss: {:.4f}, accuracy: {:.2f}'.format(test_loss / len(test_loader.dataset), 100 * correct / len(test_loader.dataset)))
profile
Yes, Code Wins Arguments!!

0개의 댓글