[PyTorch] Forward Net

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

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

input_size = 784
hidden_size = 100
output_size = 10
num_epochs = 2
batch_size = 100
learning_rate = 0.001

train_dataset = torchvision.datasets.MNIST(root='/content/drive/MyDrive/Study/data', train=True, transform=transforms.ToTensor(), download=True)

test_dataset = torchvision.datasets.MNIST(root="/content/drive/MyDrive/Study/data", train=False, transform=transforms.ToTensor())

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

examples = iter(train_loader)
samples, labels = examples.next()
samples.shape, labels.shape

(torch.Size([100, 1, 28, 28]), torch.Size([100]))

for i in range(6):
  plt.subplot(2,3,i+1)
  plt.imshow(samples[i][0])
  # plt.imshow(samples[i][0], cmap="gray")  #흑백
plt.show()

class NeuralNet(nn.Module):
  def __init__(self, input_size, hidden_size, output_size):
    super(NeuralNet, self).__init__()
    self.linear1 = nn.Linear(input_size, hidden_size)
    self.relu1 = nn.ReLU()
    self.linear2 = nn.Linear(hidden_size, output_size)
  
  def forward(self, x):
    out = self.linear1(x)
    out = self.relu1(out)
    out = self.linear2(out)
    return out

model = NeuralNet(input_size, hidden_size, output_size)

#loss and optimizer

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

n_total_steps = len(train_loader)


for epoch in range(num_epochs):
  for i, (images, labels) in enumerate(train_loader):
    # 100, 1, 28, 28
    # 100, 784
    images = images.reshape(-1, 28*28).to(device)
    labels = labels.to(device)

    outputs = model(images)
    loss = criterion(outputs, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (i+1) % 100 == 0:
      print(f"epoch = {epoch+1} / {num_epochs}, step{i+1}/{n_total_steps}, loss={loss.item():.4f}")

  

epoch = 1 / 2, step100/600, loss=0.3824
epoch = 1 / 2, step200/600, loss=0.1703
epoch = 1 / 2, step300/600, loss=0.2758
epoch = 1 / 2, step400/600, loss=0.3444
epoch = 1 / 2, step500/600, loss=0.2894
epoch = 1 / 2, step600/600, loss=0.2160
epoch = 2 / 2, step100/600, loss=0.1872
epoch = 2 / 2, step200/600, loss=0.3154
epoch = 2 / 2, step300/600, loss=0.1302
epoch = 2 / 2, step400/600, loss=0.2134
epoch = 2 / 2, step500/600, loss=0.2352
epoch = 2 / 2, step600/600, loss=0.1436

with torch.no_grad():
  n_correct = 0
  n_samples = 0
  for images, labels in test_loader:
    images = images.reshape(-1, 28*28).to(device)
    labels = labels.to(device)
    outputs = model(images)

    _, pred = torch.max(outputs,1)
    n_samples += labels.shape[0]
    n_correct += (pred == labels).sum().item()

  acc = 100.0 * n_correct / n_samples
  print(f"accuracy = {acc}")

accuracy = 95.28