#!/usr/bin/env python
# coding: utf-8

# # **2장 사전 훈련된 신경망**
# ## 3 뉴럴토크2
# - [ImageCaptioning.pytorch - Guthub](https://github.com/deep-learning-with-pytorch/ImageCaptioning.pytorch)
# 
# ## 4 Torch Hub
# - [Torch Hub Document](https://pytorch.org/docs/stable/hub.html)
# - [Torch Hub Model List](https://pytorch.org/hub/research-models)
# - [Torch Vision Github](https://github.com/pytorch/vision)
# 
# [다음과 같은 Warning 을 출력하는 경우](https://github.com/pytorch/vision/issues/2291#issuecomment-639552166)
# ```
# Did you have `libjpeg` or `libpng` installed before building `torchvision` 
# ! pip3 install auditwheel
# ```

# In[7]:


import torch
entrypoints = torch.hub.list('pytorch/vision') # , force_reload=True)
len(entrypoints), entrypoints[::20]


# In[8]:


# Resnet18 모델 호출 및 파라미터 확인
resnet18 = torch.hub.load('pytorch/vision', 'resnet18', weights='DEFAULT')
print(torch.hub.help('pytorch/vision', 'resnet18', force_reload=True))


# ## 3 MNIST
# - Convolution Model

# In[3]:


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
torch.manual_seed(4242)

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data/ch02/mnist', train=True, download=True,
       transform=transforms.Compose([
           transforms.ToTensor(),
           transforms.Normalize((0.1307,), (0.3081,))
       ])),
    batch_size=64, shuffle=True)
train_loader


# In[4]:


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Net()
model


# In[5]:


optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
    print(f'Epoch : {epoch} >>> Current loss', float(loss))

torch.save(model.state_dict(), 'data/ch02/mnist/mnist.pth')
pretrained_model = Net()
pretrained_model.load_state_dict(torch.load('data/ch02/mnist/mnist.pth'))