【Pytorch】 Pytorch튜토리얼 6 - what is torch.nn really?

03 May 2020 in Pytorch / Docker / Git

아래의 주석을 참고해서 공부하기 바랍니다. 차근히 공부해보기

WHAT IS TORCH.NN REALLY?

1. MNIST 데이터 세트에서 기본 신경망을 학습
2. 점진적에서 하나 개의 기능을 추가 할 것 -> torch.nn, torch.optim, Dataset, DataLoader
3. 처음에는 정말 코드를 복잡하게 만들고, 그것을 torch 내부의 모듈과 함수(클래스)를 이용해서 코드를 점점 쉽게 구현해 나갈 것이다.

   1. 파일 및 이미지 다운. 파일을 torch.tensor로 변환하기

from pathlib import Path        # pathlib는 파일위치 찾기, 파일 입출력에 사용하는 모듈. 과거 os모듈. https://brownbears.tistory.com/415
import requests                 # 간편한 HTTP 요청처리를 위해 사용하는 모듈

# 1. 폴더를 만들고, MNIST 데이터 다운로드 하기
DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"          # os.path.join 과 같은 느낌.
PATH.mkdir(parents=True, exist_ok=True)

URL = "http://deeplearning.net/data/mnist/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():  # os.path.join 과 같은 느낌.
        content = requests.get(URL + FILENAME).content
        (PATH / FILENAME).open("wb").write(content)



# 2. 다운한 파일의 압축을 풀고, 파일을 Load 하여, 하나의 변수에 넣는다.
import pickle                   # 파일 load하는데 많이 쓰이는 모듈
import gzip                     # 압축된 파일의 내용을(굳이 압축안 풀고) 바로 읽을 수 있게 해주는 모듈 : https://itholic.github.io/python-gzip/
with gzip.open((PATH / FILENAME).as_posix(), "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

# 3. 파일 다운로드가 잘 되었나 확인하보자. 
from matplotlib import pyplot
import numpy as np

print(x_train.shape, '\n', type(x_train))
pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")

(50000, 784) 
 <class 'numpy.ndarray'>





<matplotlib.image.AxesImage at 0x2930e3d3088>

# 4. torch.tensor를 사용할 것이기 때문에, 
import torch

x_train, y_train, x_valid, y_valid = map(torch.tensor, (x_train, y_train, x_valid, y_valid)) # https://pytorch.org/docs/stable/tensors.html#torch.Tensor
n, c = x_train.shape

print(x_train.shape)
print(y_train.min(), y_train.max())  # 0 ~ 9까지 10개의 Class가 존재한다.

torch.Size([50000, 784])
tensor(0) tensor(9)

2. torch.nn을 사용하지 않고 신경망 구현해 보기.

nn을 이용해서, 매개변수를 정의한다면, 자동으로 requires_grad = True가 된다.
하지만 아래와 같이 매개변수를 직접 정의한다면, requires_grad = True를 직접해주어야한다.
그리고, _를 사용하면 the operation is performed in-place라는 것을 의미한다.

import math

weights = torch.randn(784, 10) / math.sqrt(784)     # 여기서 가중치 초기화 방법으로 Xavier initialisation 를 사용했다.
weights.requires_grad_()                            # Defalut =-> requires_grad=True : https://pytorch.org/docs/stable/tensors.html#torch.Tensor.requires_grad_
bias = torch.zeros(10, requires_grad=True)

# 1층 Fully connected Layer를 만든다.
def log_softmax(x):
    return x.exp().log() - x.exp().sum(-1).log().unsqueeze(-1)  
    # torch.tensor의 함수(math 모듈의 함수 NO)를 잘 이용하고 있다.  
    # x.exp().log() == x
def model(xb):
    return log_softmax(xb @ weights + bias)         # @내적 연산을 의미

# 64장을 하나의 배치로 하고, Forward를 진행해 나간다.
bs = 64  # batch size

xb = x_train[0:bs]      # a mini-batch from x
preds = model(xb)       # predictions
print(preds[0])
print(preds.shape)

tensor([-2.3311, -2.3437, -1.6287, -2.1573, -2.8016, -2.4115, -2.5137, -2.1718,
        -2.4900, -2.7065], grad_fn=<SelectBackward>)
torch.Size([64, 10])

loss function을 만들어 보자

def nll_loss(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll_loss        # 함수 포인터는 이처럼 이용하면 된다.

yb = y_train[0:bs]
print(yb.shape) # torch.Size([64]) -> 64장의 이미지 각각의 class가 적혀 있다.
print(loss_func(preds, yb))

torch.Size([64])
tensor(2.3936, grad_fn=<NegBackward>)

accuracy function을 만들어보자.

def accuracy(preds_before, yb):
    # 각 예측에 대해 가장 큰 값을 가진 인덱스가 목표 값과 일치함을 판단합니다.
    preds = torch.argmax(preds_before, dim=1) 
    # dim : the dimension to reduce. If None, the argmax of the flattened input is returned.
    # preds_before -> shape : (64, 10) -> (64)
    return (preds == yb).float().mean()

    """
    print(preds) -> tensor([3, 3, 3, 6, 6, 3, 3.....  9, 3, 6])
    print(preds.shape) -> torch.Size([64])
    """
print(accuracy(preds, yb))

tensor(0.0312)

이제 훈련을 시켜보자.

loop를 통해서,
데이터 가져오기 -> forward -> loss계산 -> backward -> 가중치 갱신
이 되는 것을 확인하라.

from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    # x_train.shape == (50000,784) 
    # n == 50000 , bs = 64
    for i in range((n - 1) // bs + 1):
        # set_trace()
        # 튜토리얼 문서에 의하면, 이 코드를 디버깅하면서 한줄한줄 확인하고 싶다면 위의 주석을 풀라 했다.
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()
        # 여기서 x_train의 가장 마지막 16개의 사진은 학습에 사용 못 된다.

print(loss_func(model(xb), yb), "\n",accuracy(model(xb), yb)) # 가장 마지막 배치에 대한, loss와 accuracy를 확인해보자

tensor(0.0827, grad_fn=<NegBackward>) 
 tensor(1.)

3. Using torch.nn.functional

위에서 했던 동일한 작업을 수행하기 위해, PyTorch의 nn클래스를 활용하여보다 간결하고 유연한 코드를 만들어보자.
우선 torch.nn.functional를 사용해서 코드를 만들어 봅시다. 여기에는 torch.nn의 모든 기능이 포함되어 있습니다.

# 새로운 학습을 위해.. 다시!
xb = x_train[0:bs]      # a mini-batch from x
yb = y_train[0:bs]
weights = torch.randn(784, 10) / math.sqrt(784)     # 여기서 가중치 초기화 방법으로 Xavier initialisation 를 사용했다.
weights.requires_grad_()                            # Defalut =-> requires_grad=True : https://pytorch.org/docs/stable/tensors.html#torch.Tensor.requires_grad_
bias = torch.zeros(10, requires_grad=True)

import torch.nn.functional as F

loss_func = F.cross_entropy  # 위의 nll_loss 처럼, 함수 포인터는 이처럼 이용하면 된다.

def model(xb):
    return xb @ weights + bias

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

tensor(2.4334, grad_fn=<NllLossBackward>) tensor(0.)

4. Refactor using nn.Module

nn.Module및 nn.Parameter를 적극적으로 이용합니다.

from torch import nn
# import torch.nn as nn

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
         # nn.Parameter(텐서) : 이 텐서를 parameter로 이용할 것을 명명한다. grad를 알아서 해준다.
        self.weights = nn.Parameter(  torch.randn(784, 10) / math.sqrt(784)  )
        self.bias = nn.Parameter(torch.zeros(10))

    def forward(self, xb):
        return xb @ self.weights + self.bias

model = Mnist_Logistic()
print(loss_func(model(xb), yb)) # loss값 구하기.

tensor(2.3004, grad_fn=<NllLossBackward>)

# 위에서는 weights -= weights.grad * lr -> weights.grad.zero_() 과 같은 과정을 bias에도 반복했었지만.. 여기서는 쉽게 할 수 있다.
with torch.no_grad():
    for p in model.parameters(): 
        p -= p.grad * lr
    model.zero_grad()

from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

def fit():
    for epoch in range(epochs):
        for i in range((n - 1) // bs + 1):
            start_i = i * bs
            end_i = start_i + bs
            xb = x_train[start_i:end_i]
            yb = y_train[start_i:end_i]
            pred = model(xb)
            loss = loss_func(pred, yb)

            loss.backward()
            with torch.no_grad():
                for p in model.parameters():  p -= p.grad * lr
                model.zero_grad()

fit()

print(loss_func(model(xb), yb))

5. Refactor using nn.Linear

위에서는 weight와 bias를 직접 정의했지만,
이제는 nn.Linear를 사용해서 코드를 구현해보자

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

    def forward(self, xb):
        return self.lin(xb)

model = Mnist_Logistic()
print(loss_func(model(xb), yb))

fit()  # 위에 있는 함수 그대로 사용해도 된다.

6. Refactor using optim

위에서 했던,

with torch.no_grad():
    for p in model.parameters(): p -= p.grad * lr
    model.zero_grad()

과정을 optim를 이용해서 쉽게 구현해보자. 그냥

opt.step()
opt.zero_grad()

를 하면 된다!

from torch import optim
# 더 아래에서도 사용하기 위해서, 굳이 이렇게 get_model이라는 함수를 구현했다. 
def get_model():
    model = Mnist_Logistic()
    return model, optim.SGD(model.parameters(), lr=lr)

model, opt = get_model()
"""
model = Mnist_Logistic()
opt = optim.SGD(model.parameters(), lr=lr)
"""
print(loss_func(model(xb), yb))

lr = 0.5 
epochs = 2

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

7. Refactor using Dataset

지금까지

start_i = i * bs
end_i = start_i + bs
xb = x_train[start_i:end_i]
yb = y_train[start_i:end_i]

를 했지만, 파이토치의 an abstract Dataset class인, TensorDataset를 이용해보자. (텐서의 첫 번째 차원을 따라 반복, 인덱싱 및 슬라이스하는 방법을 제공)
TensorDataset에는 len , getitem 이라는 좋은 함수가 있다.
* 자세한 사항은 이 홈페이지를 공부하자 **언젠간 공부해야 한다!! ***

from torch.utils.data import TensorDataset
train_ds = TensorDataset(x_train, y_train)
""" 위에서는 이렇게 했었다. 
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
""" 
# 이제는 이런 식으로 구현할 것이다. 
start_i = i * bs
end_i = start_i + bs
xb,yb = train_ds[start_i : end_i] # 처음에 train_ds정의를 튜플로 했으므로, 항상 output도 튜플로 해준다.

lr = 0.5 
epochs = 2

model, opt = get_model()

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        xb, yb = train_ds[i * bs: i * bs + bs]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

8. Refactor using DataLoader

이제는 [i * bs: i * bs + bs] 이렇게 하는 것도 싫다.
DataLoader를 사용함으로써 배치 관리를 쉽게 할 수 있다.

from torch.utils.data import DataLoader

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs)
"""위에서는 이렇게 했었다. 
for i in range((n-1)//bs + 1):
    xb,yb = train_ds[i*bs : i*bs+bs]
    pred = model(xb)
"""
# 이제는 이런 식으로 구현할 것이다.
for xb,yb in train_dl:
    pred = model(xb)
    break

lr = 0.5 
epochs = 2

model, opt = get_model()
loss_func = F.cross_entropy

for epoch in range(epochs):
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

9. Add validation

앞으로 우리는 a validation set을 사용할 것이다. in order to identify if you are overfitting. 유의사항

1. train dataset은 suffling(섞거나, 랜덤하게 뽑아서) 이용했지만(오버피팅 막기 위해), validation set에서는 그런 작업이 필요없다.
2. validation에서는 train보다 2배이상의 batch사이즈를 사용할 것이다. 역전파를 하지 않기 때문에 메모리 사용이 적기 때문이다.
3. 큰 배치를 사용해서, 빠르게 Loss값을 구하고 validation 값을 확인한다.

# Train dataset
train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)
# Validation dataset
valid_ds = TensorDataset(x_valid, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=bs * 2)

lr = 0.5 
epochs = 2

model, opt = get_model()
loss_func = F.cross_entropy

# point 1 : validation (or inference)에서는 torch.no.grad() 로 처리한다
# point 2 : train, eval하기 전에, model.train() / model.eval()를 해준다. 배치Norm이나 DropOut과 같은 layer처리를 알아서 바꿔준다. torch.nn의 매소드라고 할 수 있다.

for epoch in range(epochs):
    model.train()
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

    model.eval()
    with torch.no_grad():
        # print((xb, yb) for xb, yb in valid_dl) -> <generator object <genexpr> at 0x000001A8AFF87848>
        # print(loss_func(model(xb), yb) for xb, yb in valid_dl)  -> { []를 치지 않아도, [변수_ for _ in _] 가 잘 동작한다...] } -> <generator object <genexpr> at 0x000001A89EC026C8>
        # print(list(loss_func(model(xb), yb) for xb, yb in valid_dl))  -> [tensor(0.3860), tensor(0.4615), tensor(0.4938), tensor(0.5899),  .... 
        # print(len(list(loss_func(model(xb), yb) for xb, yb in valid_dl)))   -> [79] 백터
        valid_loss = sum(loss_func(model(xb), yb) for xb, yb in valid_dl)

    print(epoch,"번째 epoch에서 valid_loss값은" ,valid_loss / len(valid_dl))

10. 지금까지 했던 것을 함수로 만들기!

loss_batch // fit // get_data 라는 이름의 함수를 만들지.

1. loss_batch : one batch에 대해서 loss를 구해주는 함수
2. fit : loss_batch함수를 이용해서, 모델 전체를 train, validation 해주는 함수.
3. get_data :

def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )

def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)  # 한 batch 즉 64장에 대한, 총 (평균) loss를 계산한다.

    if opt is not None:  # validation, inference를 위해서 만들어 놓는 옵션.
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)  # train 과정 중, 이 값은 필요 없다. validation을 위해 return을 만들어 놓았다.

import numpy as np # 맨 아래 val_loss를 구하기 위해 numpy를 잠깐 쓴다.

def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)

        model.eval()
        with torch.no_grad():
            # [*,*,*,*,*,*,*],[+,+,+,+,+,+,+] <= zip( [*,+] , [*,+] , [*,+] , [*,+] ...)
            # print( list(zip( * [ loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl])) ) 
            # print( list([* [ loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]]) )
            # 여기서 [] for in 관계가 어떻게 되는 거지? 는 위의 주석을 풀어보면 된다.
            losses, nums = zip( * [ loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl] )  
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print(epoch, val_loss)

# 자 이제 우리가 만든 함수를 돌려보자
train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
model, opt = get_model()
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

11. nn.linear말고, CNN 사용하기

import torch.nn as nn
import torch.nn.functional as F
class Mnist_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1)

    def forward(self, xb):
        xb = xb.view(-1, 1, 28, 28) # 첫 input size = [[이미지 장 수, 784]]
        xb = F.relu(self.conv1(xb))
        xb = F.relu(self.conv2(xb))
        xb = F.relu(self.conv3(xb))
        xb = F.avg_pool2d(xb, 4) # https://pytorch.org/docs/stable/nn.functional.html#pooling-functions
        return xb.view(-1, xb.size(1))

lr = 0.1

from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.optim as optim

epochs = 2
bs = 64
train_ds = TensorDataset(x_train, y_train)
valid_ds = TensorDataset(x_train, y_train)

loss_func = F.cross_entropy
train_dl, valid_dl = get_data(train_ds, valid_ds, bs)

# 하기 전 #10의 3개의 함수 선언 하기.
model = Mnist_CNN()
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.41726013900756836
1 0.27703515412330626

12. nn Sequential

nn의 하나의 클래스인 Sequential을 이용해 보자.
Sequential객체는 순차적으로 내부에 포함 된 각 모듈을 실행한다.
여기서 주의할 점은, Sequential을 정확하게 이용하기 위해서, view를 위한 layer를 하나 정의해주어야한다.
아래의 내용은 함수 포인터, 함수 input매개변수 등 다양한 사항을 고민해야한다.

class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


def preprocess(x):
    return x.view(-1, 1, 28, 28)

# http://hleecaster.com/python-lambda-function/
# lambda input매개변수 : 함수내부 연산 
# 변수명_a = lambda ~~ : ~~  ; -> 변수명_a는 함수 포인터이다.

"""
    Mnist_CNN에서는 forward에서 xb = xb.view(-1, 1, 28, 28)를 쉽게 했지만, Sequential에서는 그렇게 하지 못한다.
    따라서 이러한 처리를 하는 방법은 바로 아래 코드처럼 구현을 하는 것이다.
    Lambda(preprocess) 에서 preprocess(x_input)이 들어가면서 Sequential aclass가 나중에 실행될 것이다. 
    Sequential에 순차적으로 넣는 매개변수 하나하나는 꼭! nn.Module로 정의된 Class 이여야 한다. # 이것이 view용 레이어 처리 방법이다.

"""
model = nn.Sequential(
    Lambda(preprocess),  
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AvgPool2d(4),
    Lambda(lambda x: x.view(x.size(0), -1)), 
)
# 맨 아래에서도 preproces같은 함수 포인터가 들어가면 좋지만, 굳이 def preprocess: 하기 귀찮으므로...

lr = 0.5
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.8372452193450928
1 0.747732577419281

13. Wrapping DataLoader

모든 2d single channel image라면 input을 무조건 받을 수 있는, model을 구현해 보자.
위에서는 preprocess라는 함수를 정의하고, nn.Sequential(Lambda(preprocess), … ; 처럼 사용했었다.
그러지 말고, 아에 처음부터 train_dl, valid_dl을 view처리를 한 상태에서 model에 집어넣자.

def preprocess(x, y):
    return x.view(-1, 1, 28, 28), y                                         # 굳이 계속 none일 y를 새로 정의한 이유가 뭘까?????

class WrappedDataLoader:
    def __init__(self, dl, func):
        self.dl = dl
        self.func = func

    def __len__(self):
        return len(self.dl)

    # WrappedDataLoader클래스로 정의한 객체에서, iter함수를 사용하고 싶다면 이것을 정의한다.
    # yield == 생성기(generator) == https://python.bakyeono.net/chapter-7-4.html
    def __iter__(self): 
        batches = iter(self.dl)
        for b in batches:
            # print(type(b))                                                # list
            # print(len(b) , b[0].shape)                                    # 2(?) torch.Size([64, 784])
            # print(self.func(*b))                                          # 2장의 이미지가 view처리 되어 나온다.
            # print(len(self.func(*b)), self.func(*b)[0].shape)             # 2, torch.Size([64, 1, 28, 28])
            yield (self.func(*b))  # preprocess == func으로 동작한다. *b가 들어가는 것은.. 잘 모르겠다. 어째서지?

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

# 위에서 train_dl을 잘 만져놓았으므로, model의 input은 무조건 적절한 사이즈의 input일 것이다.
model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    Lambda(lambda x: x.view(x.size(0), -1)),
)

opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.34880866147994993
1 0.2753501031970978

14. GPU사용해서 가속하기

print(torch.cuda.is_available())

dev = torch.device("cuda" if torch.cuda.is_available() else torch.device("cpu") )
# dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")  <- 홈페이지 방법. 둘다 된다.

def preprocess(x, y):
    return x.view(-1, 1, 28, 28).to(dev), y.to(dev)  # 핵심 포인트!!


train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

model.to(dev)                                        # 이것도 핵심 포인트!!!
                                                     # 4Classifier에서 net.to(dev)를 검색해서 보자. <- 똑같은 방법 사용!
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
fit(epochs, model, loss_func, opt, train_dl, valid_dl)