from tqdm import tqdm
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import torch.nn as nn
from torchsummary import summary
from torch.optim import SGD
import torch.nn as nn
BATCH_SIZE = 32
LR = 0.1
DEVICE = 'cuda'
EPOCHS = 20
N_TRAIN_SAMPLES = 60000
DEVICE = 'cuda'
: NVIDIA GPU에서 돌린다train_ds = MNIST(root='.', train=True, transform=ToTensor(), download=True)
train_loader = DataLoader(dataset=train_ds, batch_size=BATCH_SIZE, shuffle=True)
n_train_samples = len(train_ds) # 60000
train_loader
: 전체 데이터를 배치만큼 뽑아주는 친구class MLP(nn.Module):
def __init__(self):
super(MLP, self).__init__()
self.fc1 = nn.Linear(in_features=784, out_features=128) #784개 input layer, 출력 128개
self.fc1_act = nn.Sigmoid()
self.fc2 = nn.Linear(in_features=128, out_features=64) #첫 레이어 128개 뉴런, 출력 64개
self.fc2_act = nn.Sigmoid()
self.fc3 = nn.Linear(in_features=64, out_features=10) #64개 뉴런, 출력 10개
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
x = self.fc1(x)
x = self.fc1_act(x)
x = self.fc2(x)
x = self.fc2_act(x)
x = self.fc3(x)
x = self.softmax(x)
return x
model = MLP().to(DEVICE)
loss_fn = nn.CrossEntropyLoss() #loss function
optimizer = SGD(model.parameters(), lr=LR) #w:=w-lr*dJ/dw
for epoch in range(EPOCHS):
loss_epoch = 0.
for imgs, labels in tqdm(train_loader):
#backward propagation
imgs = imgs.to(DEVICE).reshape(imgs.shape[0], -1)
labels = labels.to(DEVICE)
preds = model(imgs)
loss = loss_fn(preds, labels)
#backward propagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_epoch += loss.item() * imgs.shape[0]
loss_epoch /= n_train_samples
print(f"\nEpoch: {epoch + 1}")
print(f"Train Loss: {loss_epoch:.4f}\n")
for imgs, labels in tqdm(train_loader)
: train_loader
에서 하나씩 뽑아서 아래 코드를 실행한다.
imgs.to(DEVICE).reshape(imgs.shape[0], -1)
: GPU에 업로드하고, (배치사이즈, 784)의 행렬로 바꿔준다.
labels.to(DEVICE)
: 정답값도 GPU에서 연산시키기 위해 GPU에 업로드
model(imgs)
: 예측 모델 만들기
- 여기서 model.forward에 정의했던 연산들이 수행된다
loss_fn(preds, labels)
: cross entropy error에서 로스를 통해 배치들의 loss의 평균을 구한다
optimizer.zero_grad()
:그냥 무조건 해준다고 생각하기(gradient의 누적을 막기 위해)->일단은 backward하기 전에 해준다고 생각하기
loss.backward()
: backpropagation 미분들을 다 구해준다. dJ/dw가 다 구해짐
optimizer.step()
: 모든 파라미터에 대해서 업데이트를 해줌 w:=w-al*dJ/dw
loss_epoch /= n_train_samples
: 에폭의 평균 loss값