%matplotlib inline

"""
Explanation: Minist retrainning(finetuning)
Minist 예제에서 Network를 구조를 마지막 fully connected (linear) 부분의 차원을 수정한 다음 재훈련을 하도록 하겠습니다.
이것을 수행하기 위해서는 다음과 같은 방법이 있는데, finetuning에 해당하는  방법3을 구현하도록 하겠습니다.

방법1 : Minist Network를 새로 작성한다.


python
self.fc1 = nn.Linear(64*7*7, 1024)
self.fc2 = nn.Linear(1024, 10)

위의 부분을 아래와 같이 수정하고

python
self.fc1 = nn.Linear(64*7*7, 512)
self.fc2 = nn.Linear(512, 10)

그리고 trainning을 다시 수행합니다.

방법2 : 기존의 Minist Network를 불러들인 후, 필요부분을 수정함.


python
model = MnistModel()
model.fc1 = nn.Linear(64*7*7, 512)
model.fc2 = nn.Linear(512, 10)

그리고 trainning을 다시 수행합니다.

방법3 :  기존의 Minist Network와 사전학습된(pretrained) parameter를 복원한 후, Network를 수정하고 수정부분만을 trainning 함. 
변경된 network의 parameter만을 update하므로 방법2보다 빨리 training을 할 수 있고, 정확도도 유사합니다.


```python
model = MnistModel()

load parameter of MnistModel
checkpoint = torch.load(checkpoint_filename)
model.load_state_dict(checkpoint)
modify last two layer in model = MnistModel()
model.fc1 = nn.Linear(6477, 512)
model.fc2 = nn.Linear(512, 10)
specify parameters to update
fc_parameters = [
    {'params': model.fc1.parameters()},
    {'params': model.fc2.parameters()}
]
assign parameter in optimizer
optimizer = torch.optim.Adam(fc_parameters, lr=0.0001)
```
End of explanation
"""


import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np

is_cuda = torch.cuda.is_available() # cuda 사용가능시, True
checkpoint_filename = 'minist.ckpt'

batch_size = 50
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data', train=True, download=True, transform=transforms.ToTensor()),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data', train=True, transform=transforms.ToTensor()),
    batch_size=100)


"""
Explanation: 1. 입력DataLoader 설정

train 데이터로 loader를 지정 (dataset은 Minist, batch 사이즈 50, shuffle를 실행)
test 데이터로 loader를 지정 (dataset은 Minist, batch 사이즈 1000)
End of explanation
"""


class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()
        # input is 28x28
        # padding=2 for same padding
        self.conv1 = nn.Conv2d(1, 32, 5, padding=2)
        # feature map size is 14*14 by pooling
        # padding=2 for same padding
        self.conv2 = nn.Conv2d(32, 64, 5, padding=2)
        # feature map size is 7*7 by pooling
        self.fc1 = nn.Linear(64*7*7, 1024)
        self.fc2 = nn.Linear(1024, 10)
        
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 64*7*7)   # reshape Variable
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)
    
model = MnistModel()
#load parameter of MnistModel
checkpoint = torch.load(checkpoint_filename)
model.load_state_dict(checkpoint)

### don't update model parameters
for param in model.parameters() :
    param.requires_grad = False
#modify last two layer in model = MnistModel()
model.fc1 = nn.Linear(64*7*7, 512)
model.fc2 = nn.Linear(512, 10)

fc_parameters = [
    {'params': model.fc1.parameters()},
    {'params': model.fc2.parameters()}
]
optimizer = torch.optim.Adam(fc_parameters, lr=0.0001)

if is_cuda :  model.cuda()
loss_fn = nn.NLLLoss() 

"""
Explanation: 2. 사전 설정
* model
* loss
* opimizer
End of explanation
"""


# trainning
model.train()
train_loss = []
train_accu = []
for epoch in range(3):
    for i, (image, target) in enumerate(train_loader):
        if is_cuda :  image, target = image.cuda(), target.cuda() 
        image, target = Variable(image), Variable(target) # 입력image Target 설정 
        output = model(image) # model 생성
        loss = loss_fn(output, target) #loss 생성
        optimizer.zero_grad() # zero_grad
        loss.backward() # calc backward grad
        optimizer.step() # update parameter
        
        pred = output.data.max(1)[1]
        accuracy = pred.eq(target.data).sum()/batch_size
        
        train_loss.append(loss.data[0])
        train_accu.append(accuracy)

        if i % 300 == 0:
            print(i, loss.data[0])

plt.plot(train_accu)

plt.plot(train_loss)

"""
Explanation: 3. Trainning loop
* (입력 생성)
* model 생성
* loss 생성
* zeroGrad
* backpropagation
* optimizer step (update model parameter)

```
End of explanation
"""


model.eval()
correct = 0
for image, target in test_loader:
    if is_cuda :  image, target = image.cuda(), target.cuda() 
    image, target = Variable(image, volatile=True), Variable(target)
    output = model(image)
    prediction = output.data.max(1)[1]
    correct += prediction.eq(target.data).sum()

print('\nTest set: Accuracy: {:.2f}%'.format(100. * correct / len(test_loader.dataset)))

"""
Explanation: 4. Predict & Evaluate
End of explanation
"""


