import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

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

import torch.utils

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

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

import numpy as np

plt.figure(figsize=(20,5))
for i in range(20):
    img=imgs[i]
    #进行轴变换         
    np_img = img.numpy().transpose((1, 2, 0))
    plt.subplot(2,10,i+1)
    plt.imshow(np_img,cmap=plt.cm.binary)
    plt.axis("off")

import torch.nn.functional as F
num_classes=10
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(3,64,kernel_size=3)
        self.pool1=nn.MaxPool2d(kernel_size=2)
        self.conv2=nn.Conv2d(64,64,kernel_size=3)
        self.pool2=nn.MaxPool2d(kernel_size=2)
        self.conv3=nn.Conv2d(64,128,kernel_size=3)
        self.pool3=nn.MaxPool2d(kernel_size=2)
        
        self.fc1=nn.Linear(512,256)
        self.fc2=nn.Linear(256,num_classes)
    def forward(self,x):
        x=self.pool1(F.relu(self.conv1(x)))
        x=self.pool2(F.relu(self.conv2(x)))
        x=self.pool3(F.relu(self.conv3(x)))
        x=torch.flatten(x,start_dim=1)
        x=F.relu(self.fc1(x))
        x=self.fc2(x)
        return x

from torchinfo import summary

model=Model().to(device)
summary(model)

loss_fn=nn.CrossEntropyLoss()
learning_rate=1e-2
opt=torch.optim.SGD(model.parameters(),lr=learning_rate)

def train(dataloader,model,loss_fn,optimizer):
    size=len(dataloader.dataset)
    num_batches=len(dataloader)
    train_loss,train_acc=0,0
    for X,y in dataloader:
        X,y=X.to(device),y.to(device)
        pred=model(X)
        loss=loss_fn(pred,y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        train_acc+=(pred.argmax(1)==y).type(torch.float).sum().item()
        train_loss+=loss.item()
    train_acc/=size
    train_loss/=num_batches
    return train_acc,train_loss

def test(dataloader,model,loss_fn):
    size=len(dataloader.dataset)
    num_batches=len(dataloader)
    test_loss,test_acc=0,0
    with torch.no_grad():
        for imgs,target in dataloader:
            imgs,target=imgs.to(device),target.to(device)
            target_pred=model(imgs)
            loss=loss_fn(target_pred,target)
            test_loss+=loss.item()
            test_acc+=(target_pred.argmax(1)==target).type(torch.float).sum().item()
    test_acc/=size
    test_loss/=num_batches
    return test_acc,test_loss

epochs=10
train_loss=[]
train_acc=[]
test_loss=[]
test_acc=[]
for epoch in range(epochs):
    model.train()
    epoch_train_acc,epoch_train_loss=train(train_dl,model,loss_fn,opt)
    model.eval()
    epoch_test_acc,epoch_test_loss=test(test_dl,model,loss_fn)
    train_loss.append(epoch_train_loss)
    train_acc.append(epoch_train_acc)
    test_loss.append(epoch_test_loss)
    test_acc.append(epoch_test_acc)
    template=('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%，Test_loss:{:.3f}')
    print(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss, epoch_test_acc*100, epoch_test_loss))
print('Done')
        

import matplotlib.pyplot as plt
#隐藏警告
import warnings
warnings.filterwarnings("ignore")              
plt.rcParams['font.sans-serif']    = ['SimHei'] 
plt.rcParams['axes.unicode_minus'] = False     
plt.rcParams['figure.dpi']         = 100        

epochs_range = range(epochs)

plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)

plt.plot(epochs_range, train_acc, label='Training Accuracy')
plt.plot(epochs_range, test_acc, label='Test Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss, label='Training Loss')
plt.plot(epochs_range, test_loss, label='Test Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()