[PyTorch] 03. torch.nn 모듈 조립으로 CNN 만들기

import torchvision.datasets as datasets

train_data = datasets.MNIST(root='./data', train=True, download=True)
test_data = datasets.MNIST(root='./data', train=False, download=True)

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1_conv = nn.Conv2d(1, 6, 5, 1) # 입력 1개, 출력 6개, 필터 크기는 5x5, 1칸 단위로 이동하면서 필터를 씌운다
        self.layer1_relu = nn.ReLU()             # 활성화 함수. ReLU(x) 는 max(x, 0)과 같다
        self.layer1_pool = nn.MaxPool2d(2)       # 각 2x2 칸마다 최대값 하나씩만을 남긴다
        self.layer2_conv = nn.Conv2d(6, 16, 5, 1)
        self.layer2_relu = nn.ReLU()
        self.layer2_pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16*4*4, 10)
    
    def forward(self, x):
        x1 = self.layer1_conv(x)  # 1x28x28 형식의 데이터가 6x24x24 형식으로 변환된다
        x2 = self.layer1_relu(x1) # 
        x3 = self.layer1_pool(x2) # 6x24x24 형식의 데이터가 6x12x12 형식으로 변환된다
        x4 = self.layer2_conv(x3) # 6x12x12 형식의 데이터가 16x8x8 형식으로 변환된다
        x5 = self.layer2_relu(x4) # 
        x6 = self.layer2_pool(x5) # 16x8x8 형식의 데이터가 16x4x4 형식으로 변환된다
        x7 = x6.view(-1, 256)     # 16x4x4 형식의 데이터가 256-벡터로 변환된다
        x8 = self.fc(x7)          # 256-벡터가 10-벡터로 변환된다
        return x8

net = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters())

# TypeError: conv2d(): argument 'input' (position 1) must be Tensor, not Image
net(train_data[0][0]) 

# 이제 _data의 각 요소는 (PIL Image, Tensor) 대신 (Tensor, Tensor) 가 된다.
train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())
test_data = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())

# * transform=transforms.ToTensor()로 이미 변환된 데이터라고 가정
# 모델에 1x28x28 입력이 아닌 (batch_size)x1x28x28 입력을 전해줘야 한다.

x_input = train_data[0][0] # x_input은 1x28x28 크기의 텐서
net(x_input[None,:])       # [None,:] 으로 축을 추가해 텐서 크기를 1x1x28x28로 만든다

criterion = nn.CrossEntropyLoss()

train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())
test_data = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())


class MyModel(nn.Module):
    """MyModel 클래스 정의는 위와 동일"""
	pass

net = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters())

# training for 1 epoch
st = time.time()
print(f"training with {len(train_data)} data... ", end='')
for epoch in range(1):
    for x,y in train_data:
        y_pred = net(x[None,:]) # batch를 지정하지 않았으므로 이전과 같이 [None,:]으로 변환
        y = y[None]             # 라벨(y)도 마찬가지로 크기 1인 batch로 변환한다
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
print(f"done (elapsed {time.time() - st}s)")

# test accuracy
with torch.no_grad(): # 테스트 시에는 gradient를 계산할 필요가 없다
    acc, tot = 0, 0   # 맞은 개수, 전체 개수
    for x, y in test_data:
        y_pred = net(x[None,:]).argmax() # argmax를 이용해 가장 유력한 후보를 답으로 제시
        tot += 1
        acc += int(y_pred==y)
    print(f"test accuracy : {acc}/{tot} ({100*acc/tot}%)")

train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())
test_data = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())

train_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=50)
test_loader = torch.utils.data.DataLoader(test_data)

st = time.time()
print(f"training with {len(train_data)} data... ", end='')
for epoch in range(1):
    for x,y in train_loader:
        y_pred = net(x)      # 이젠 batch data이므로 입력에 특별한 처리가 필요하지 않다
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
print(f"done (elapsed {time.time() - st}s)")

# test accuracy
with torch.no_grad():
    acc, tot = 0, 0
    for x, y in test_loader:
        y_pred = net(x)
        # 이전과 달리 batch_size가 1이 아니므로 argmax(dim=1) 설정 필요
        # 자세한 설명은 https://pytorch.org/docs/stable/torch.html#torch.argmax 참조
        acc += (y==y_pred.argmax(1)).sum()
        tot += len(y)
    print(f"test accuracy : {acc}/{tot} ({100*acc/tot}%)")