import argparse
import pandas as pd
import torchvision.datasets
from torch.utils.data import DataLoader,Dataset
import torch
import torch.nn as nn
from torch import optim
from torch.utils.tensorboard import SummaryWriter
from scipy.io import loadmat
import numpy as np


#可以在终端中输入指令训练测试
parser=argparse.ArgumentParser(description="my code")
parser.add_argument('--batch_size',type=int,default=64,help='一批样本个数')
parser.add_argument('--lr',type=float,default=0.001,help='学习率')
parser.add_argument('--epoch',type=int,default=10,help='训练次数')
parser.add_argument('--cuda',action='store_true',help="是否显卡")
opt=parser.parse_args()

if opt.cuda:
    device=torch.device("cuda")
else:device=torch.device("cpu")


train_data=torchvision.datasets.CIFAR10(root="C:\\Users\Anyzhuguai\PycharmProjects\pythonProject\\venv\data",transform=torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(0.5,0.5,0.5)]),download=False,train=True)
test_data=torchvision.datasets.CIFAR10(root="C:\\Users\Anyzhuguai\PycharmProjects\pythonProject\\venv\data",transform=torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(0.5,0.5,0.5)]),download=False,train=False)

train_set=DataLoader(train_data,batch_size=opt.batch_size,shuffle=True)
test_set=DataLoader(test_data,batch_size=opt.batch_size,shuffle=True)

#创建模型
class Trainer (nn.Module):
    def __init__(self):
        super(Trainer,self).__init__()
        self.flatten=nn.Flatten()
        self.conv1=nn.Conv2d(3,32,3,stride=1,padding=0)
        self.conv2 = nn.Conv2d(32, 16, 3, stride=1, padding=0)
        self.conv3=nn.Conv2d(16,1,3,stride=1,padding=0)
        self.pool = nn.MaxPool2d(kernel_size=3, stride=1)
        self.Linear1=nn.Linear(20*20,10)


    def forward(self,x):
        x=self.pool(torch.relu(self.conv1(x)))
        x=self.pool(torch.relu(self.conv2(x)))
        x=self.pool(torch.relu(self.conv3(x)))
        x=self.flatten(x)
        x=self.Linear1(x)

        return x

#函数实例化，创建损失函数，优化器
trainer=Trainer()
trainer=trainer.to(device)
loss=nn.CrossEntropyLoss()
optimizer=optim.Adam(trainer.parameters(),lr=opt.lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
writer=SummaryWriter(log_dir="C:\\Users\Anyzhuguai\PycharmProjects\pythonProject\\venv\\venv\logs")


#开始训练
i=0
for i in range(opt.epoch):
    i=i+1
    running_loss=0.0
    for data in train_set:
        imgs,targets=data
        optimizer.zero_grad()
        imgs,targets=imgs.to(device),targets.to(device)
        ouput=trainer(imgs)
        los=loss(ouput,targets)
        los.backward()
        optimizer.step()

        running_loss=running_loss+los.item()
    scheduler.step()
    writer.add_scalar(tag="trian_loss",scalar_value=running_loss,global_step=i)
    print("第{}训练,损失值是{}".format(i,running_loss))
writer.close()
print("训练结束")

#开始测试
correct=0
total=0
with torch.no_grad():
    for data in test_set:
        imgs,targets=data
        imgs, targets = imgs.to(device), targets.to(device)
        output = trainer(imgs)
        predicted=torch.argmax(output,dim=1)
        total=total+targets.size(0)
        correct=correct+(predicted==targets).sum().item()

accuracy=correct/total

print('测试集准确率{:.2%}'.format(accuracy))