# 姓   名：熊灿
# 开发时间：2021/11/25 14:44
import torch
import torchvision.datasets
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
import torch.nn as nn

import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

# 调用torchvision中的models中的resnet网络结构
import torchvision.models as models

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

cuda_is_available = torch.cuda.is_available()

transform_train = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
])
transforms_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
])
# 超参数
Download = True
EPOCHS = 10
LR = 0.001
BATCHSIZE = 64

train_data = torchvision.datasets.CIFAR10(root="./dataset",train=True,transform=transform_train,download=Download)
test_data = torchvision.datasets.CIFAR10(root="./dataset",train=False,transform=transforms_test,download=Download)

train_size = len(train_data)
test_size = len(test_data)


train_dataloader = DataLoader(dataset=train_data,batch_size=BATCHSIZE,shuffle=True)
test_dataloader = DataLoader(dataset=test_data,batch_size=BATCHSIZE,shuffle=True)

# 调整已经训练好的ResNet网络

# 1.调用模型
ResNet = models.resnet18(pretrained=True)
# 2.提取fc层中固定的参数
fc_features = ResNet.fc.in_features
# 3.修改输出的类别为10
ResNet.fc = nn.Linear(fc_features,10)


# 调整参数后，加载部分参数

model_dict = ResNet.state_dict()

# 1.filter out unnecessary keys
pretrained_dict = {k: v for k, v in model_dict.items() if k in model_dict}
# 2.overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
ResNet.load_state_dict(model_dict)
print(ResNet)

# 优化器选择Adam
optimizer = torch.optim.SGD(ResNet.parameters(), lr=LR)
# 损失函数
loss_func = nn.CrossEntropyLoss()  # 目标标签是one-hotted

if cuda_is_available:
    loss_func.cuda()
    ResNet.cuda()
# 开始训练

train_loss = []
train_acc = []
test_loss = []
test_acc = []
for epoch in range(EPOCHS):
    print("第{}轮训练开始:".format(epoch+1))
    total_train_loss = 0
    total_test_loss = 0
    total_train_acc = 0
    total_test_acc = 0
    for step, (b_x, b_y) in enumerate(train_dataloader):  # 分配batch data
        if cuda_is_available:
            b_x = b_x.cuda()
            b_y = b_y.cuda()
        output = ResNet(b_x)  # 先将数据放到cnn中计算output
        loss = loss_func(output, b_y)  # 输出和真实标签的loss，二者位置不可颠倒

        total_train_loss += loss.data.item()
        optimizer.zero_grad()  # 清除之前学到的梯度的参数
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 应用梯度
        _,train_pred_y = torch.max(output, 1)
        train_accuracy = (train_pred_y == b_y).sum()
        total_train_acc += train_accuracy.item()

    print("整体训练集上的loss：{}，正确率：{:.2f}%".format(total_train_loss / train_size,100 * total_train_acc / train_size))
    train_loss.append(total_train_loss / train_size)
    train_acc.append(100 * total_train_acc / train_size)

    for test_step,(test_x, test_y) in enumerate(test_dataloader):
        if cuda_is_available:
            test_x = test_x.cuda()
            test_y = test_y.cuda()
        test_output = ResNet(test_x)
        _,pred_y = torch.max(test_output, 1)
        test_loss_1 = loss_func(test_output,test_y)
        total_test_loss += test_loss_1.data.item()
        accuracy = (pred_y == test_y).sum()
        total_test_acc += accuracy.item()
                # print('| test loss: %.4f' % test_loss.data.numpy(), '| test accuracy: %.2f' % accuracy)
    print("整体测试集上的loss：{}，正确率：{:.2f}%".format(total_test_loss / test_size,100 * total_test_acc / test_size))
    test_loss.append(total_test_loss / test_size)
    test_acc.append(100 * total_test_acc / test_size)
# torch.save(cnn.state_dict(), 'cnn2.pkl')#保存模型

# 加载模型，调用时需将前面训练及保存模型的代码注释掉，否则会再训练一遍
# cnn.load_state_dict(torch.load('cnn2.pkl'))
# cnn.eval()


x = range(EPOCHS)
ax = plt.gca()
plt.plot(x,train_loss,'b',label="Train_loss")
plt.plot(x,test_loss,'r',label="test_loss")
plt.title("Train and Test loss")
plt.legend(loc='upper left')
plt.figure()

plt.plot(x,train_acc,'b',label="Train_accuracy")
plt.plot(x,test_acc,'r',label="Test_accuracy")
plt.title("Train and Test accuracy")
plt.legend(loc="lower right")
plt.show()