PyTorch-실습 : 딥러닝 모델 프로그램 만들기(1)

import torch
from torchvision import datasets, transforms
import torch.nn as nn
import torch.optim as optim

class MnistModel(nn.Module):

    def __init__(self, input_size, output_size):
        self.input_size = input_size
        self.output_size = output_size
        super().__init__()
        # 오버라이드 : 자식 클래스에서 부모클래스를 사용하게 한다.
        self.layers = nn.Sequential(
            nn.Linear(input_size, 500), nn.LeakyReLU(), nn.BatchNorm1d(500),
            nn.Linear(500, 300), nn.LeakyReLU(), nn.BatchNorm1d(300),
            nn.Linear(300, 100), nn.LeakyReLU(), nn.BatchNorm1d(100),
            nn.Linear(100, 50), nn.LeakyReLU(), nn.BatchNorm1d(50),
            nn.Linear(50, output_size),
            nn.LogSoftmax(dim=-1))

    def forward(self, x):
        y = self.layers(x)
        return y

dataset = datasets.MNIST('../data', train=True, download=True,
                         transform=transforms.Compose(transforms.ToTensor()))
x_data = dataset.data.float() / 255.
y_data = dataset.targets
x_data = x_data.view(x_data.size(0), -1)
train_cnt = int(x_data.size(0) * 0.8)
test_cnt = x_data.size(0) - train_cnt
indices = torch.randperm(x_data.size(0))
x_data = torch.index_select(x_data,dim=0,index=indices).split([train_cnt, test_cnt], dim=0)
y_data = torch.index_select(y_data,dim=0,index=indices).split([train_cnt, test_cnt], dim=0)

model = MnistModel(784, 10)
optimizer = optim.Adam(model.parameters())
loss = nn.NLLLoss()

batch_size = 128
indices = torch.randperm(x_data[0].size(0))
x = torch.index_select(x_data[0], dim=0, index=indices).split(batch_size, dim=0)
y = torch.index_select(y_data[0], dim=0, index=indices).split(batch_size, dim=0)
train_loader = zip(x, y)
total_loss = 0
model.train()
for i, (data, target) in enumerate(train_loader):
    train_model = model(data)
    train_loss = loss(train_model, target.squeeze())
    optimizer.zero_grad()
    train_loss.backward()
    optimizer.step()
    print(f"Train_{i} / loss : {float(train_loss)}")
    total_loss += float(train_loss)
print("Train_loss", total_loss / len(x))

model.eval()
indices = torch.randperm(x_data[1].size(0))
x = torch.index_select(x_data[1], dim=0, index=indices).split(batch_size, dim=0)
y = torch.index_select(y_data[1], dim=0, index=indices).split(batch_size, dim=0)
test_loader = zip(x, y)
with torch.no_grad():
    total_loss = 0
    for i, (data, target) in enumerate(test_loader):
        test_model = model(data)
        test_loss = loss(test_model, target.squeeze())
        print(f"test_{i} / loss : {float(test_loss)}")
        total_loss += float(test_loss)
print("test_loss", total_loss / len(x))

전체코드

import torch
from torchvision import datasets, transforms
import torch.nn as nn
import torch.optim as optim

class MnistModel(nn.Module):

    def __init__(self, input_size, output_size):
        self.input_size = input_size
        self.output_size = output_size
        super().__init__()
        # 오버라이드 : 자식 클래스에서 부모클래스를 사용하게 한다.
        self.layers = nn.Sequential(
            nn.Linear(input_size, 500), nn.LeakyReLU(), nn.BatchNorm1d(500),
            nn.Linear(500, 300), nn.LeakyReLU(), nn.BatchNorm1d(300),
            nn.Linear(300, 100), nn.LeakyReLU(), nn.BatchNorm1d(100),
            nn.Linear(100, 50), nn.LeakyReLU(), nn.BatchNorm1d(50),
            nn.Linear(50, output_size),
            nn.LogSoftmax(dim=-1))

    def forward(self, x):
        y = self.layers(x)
        return y


dataset = datasets.MNIST('../data', train=True, download=True,
                         transform=transforms.Compose(transforms.ToTensor()))
x_data = dataset.data.float() / 255.
y_data = dataset.targets
x_data = x_data.view(x_data.size(0), -1)
train_cnt = int(x_data.size(0) * 0.8)
test_cnt = x_data.size(0) - train_cnt
indices = torch.randperm(x_data.size(0))
x_data = torch.index_select(x_data,dim=0,index=indices).split([train_cnt, test_cnt], dim=0)
y_data = torch.index_select(y_data,dim=0,index=indices).split([train_cnt, test_cnt], dim=0)

model = MnistModel(784, 10)
optimizer = optim.Adam(model.parameters())
loss = nn.NLLLoss()

batch_size = 128
indices = torch.randperm(x_data[0].size(0))
x = torch.index_select(x_data[0], dim=0, index=indices).split(batch_size, dim=0)
y = torch.index_select(y_data[0], dim=0, index=indices).split(batch_size, dim=0)
train_loader = zip(x, y)
total_loss = 0
model.train()
for i, (data, target) in enumerate(train_loader):
    train_model = model(data)
    train_loss = loss(train_model, target.squeeze())
    optimizer.zero_grad()
    train_loss.backward()
    optimizer.step()
    print(f"Train_{i} / loss : {float(train_loss)}")
    total_loss += float(train_loss)
print("Train_loss", total_loss / len(x))


model.eval()
indices = torch.randperm(x_data[1].size(0))
x = torch.index_select(x_data[1], dim=0, index=indices).split(batch_size, dim=0)
y = torch.index_select(y_data[1], dim=0, index=indices).split(batch_size, dim=0)
test_loader = zip(x, y)
with torch.no_grad():
    total_loss = 0
    for i, (data, target) in enumerate(test_loader):
        test_model = model(data)
        test_loss = loss(test_model, target.squeeze())
        print(f"test_{i} / loss : {float(test_loss)}")
        total_loss += float(test_loss)
print("test_loss", total_loss / len(x))​