import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import torchvision
from torchvision import transforms
import os

base_dir = './四種天氣圖片數(shù)據(jù)集/'
train_dir = os.path.join(base_dir, 'train')
test_dir = os.path.join(base_dir, 'test')

transform = transforms.Compose([
    transforms.Resize((96, 96)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

train_ds = torchvision.datasets.ImageFolder(train_dir, transform=transform)
test_ds = torchvision.datasets.ImageFolder(test_dir, transform=transform)

train_ds.class_to_idx

train_dl = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)
test_dl = torch.utils.data.DataLoader(test_ds, batch_size=64, shuffle=True)

imgs, labels = next(iter(train_dl))
imgs.shape

labels

img = imgs[0]
print(img.min(), img.max())
img = (img + 1) / 2
img = img.permute(1, 2, 0)
print(img.min(), img.max())
plt.imshow(img)

# 定義模型
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)   # in: 3*96*96 ==> out: 16*94* 94
        self.bn1 = nn.BatchNorm2d(16)
        self.pool = nn.MaxPool2d(2, 2)   # out: 16*47*47
        self.conv2 = nn.Conv2d(16, 32, 3)  # out: 32*45*45 ==> pooling: 32*22*22
        self.bn2 = nn.BatchNorm2d(32)
        self.conv3 = nn.Conv2d(32, 64, 3)  # in: 32*22*22 ==> out: 64*20*20 ==> pooling: 64*10*10
        self.bn3 = nn.BatchNorm2d(64)
        self.dropout = nn.Dropout()

        # batch, channel, height, width, 
        self.fc1 = nn.Linear(64*10*10, 1024)
        self.bn_fc1 = nn.BatchNorm1d(1024)
        self.fc2 = nn.Linear(1024, 256)
        self.bn_fc2 = nn.BatchNorm1d(256)
        self.fc3 = nn.Linear(256, 4)   # 四種天氣類別

    def forward(self, x):
        x = self.pool( F.selu( self.conv1(x) ) )
        x = self.bn1(x)
        x = self.pool( F.selu( self.conv2(x) ) )
        x = self.bn2(x)
        x = self.pool( F.selu( self.conv3(x) ) )
        x = self.bn3(x)
        
        # x.view(-1, 64*10*10)
        x = nn.Flatten()(x)
        x = F.selu( self.fc1(x) )
        x = self.bn_fc1(x)
        x = self.dropout(x)
        x = F.selu( self.fc2(x) )
        x = self.bn_fc2(x)
        x = self.dropout(x)
        x = self.fc3(x)
        return x

model = Net()
if torch.cuda.is_available():
    model.to('cuda')
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()

def fit(epoch, model, train_loader, test_loader):
    correct = 0
    total = 0
    running_loss = 0
    
    # 因為bn和dropout需要手動指定訓(xùn)練模式和推理模式
    model.train()
    for x, y in train_loader:
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()
            
    epoch_loss = running_loss / len(train_loader.dataset)
    epoch_acc = correct / total
        
    # 測試過程
    test_correct = 0
    test_total = 0
    test_running_loss = 0
    model.eval()
    with torch.no_grad():
        for x, y in test_loader:
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)
            test_correct += (y_pred == y).sum().item()
            test_total += y.size(0)
            test_running_loss += loss.item()
    test_epoch_loss = test_running_loss / len(test_loader.dataset)
    test_epoch_acc = test_correct / test_total

    print('epoch:', epoch, '  loss:', round(epoch_loss, 3), '  accuracy:', round(epoch_acc, 3),
         '  test_loss:', round(test_epoch_loss, 3), '  test_accuracy:', round(test_epoch_acc, 3))
    return epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc

epochs = 10
train_loss = []
train_acc = []
test_loss = []
test_acc = []
for epoch in range(epochs):
    epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc = fit(epoch, model, train_dl, test_dl)
    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(test_epoch_loss)
    test_acc.append(test_epoch_acc)

plt.plot(range(1, epochs+1), train_loss, label='train_loss')
plt.plot(range(1, epochs+1), test_loss, label='test_loss')
plt.legend()