import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F

# 查看是否能使用GPU

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


transform = transforms.Compose([
	# 数据增强：

	# 增加随机水平旋转、垂直旋转、角度旋转、色度、亮度、饱和度、对比度的变化，以及随机的灰度化

	# transforms.RandomHorizontalFlip(),

	# transforms.RandomVerticalFlip(),

	# transforms.RandomRotation(degrees=(0, 60)),

	# transforms.ColorJitter(brightness=(0.7,1.3),contrast=(0.7,1.3),saturation=(0.7,1.3),hue=(-0.2,0.2)),

	# transforms.RandomGrayscale(p=0.1),

	# 调整图片尺寸，以适应resnet18的输入大小

	transforms.Resize((224, 224)),
	transforms.ToTensor(),
	# 数据归一化：因为torchvision数据集的输出是PILImage格式, 数据域在[0, 1]. 我们将其转换为标准数据域[-1, 1]的张量格式

	transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,
                                          shuffle=True, num_workers=48)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=32,
                                         shuffle=False, num_workers=48)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# model = torchvision.models.resnet34().to(device)

# model = torchvision.models.vgg16(pretrained=False, progress=True, num_classes=10).to(device)

model = torchvision.models.resnet18(pretrained=True).to(device)
num_ftrs = model.fc.in_features

# 调整模型output,因为只有10分类

model.fc = nn.Linear(num_ftrs, 10).to(device)

# 定义优化器

import torch.optim as optim
criterion = nn.CrossEntropyLoss()
# 梯度下降法

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9, nesterov=True)
# 自适应学习率：Adam

# optimizer = optim.Adam(model.parameters(), lr=0.001)

# optimizer = torch.optim.RMSprop(model.parameters(), lr=0.001,momentum=0.9)

for epoch in range(20):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # data中包含输入图像张量inputs, 标签张量labels

        inputs, labels = data[0].to(device), data[1].to(device)
        # 首先将优化器梯度归零

        optimizer.zero_grad()
        # 输入图像张量进网络, 得到输出张量outputs

        predicts = model(inputs)
        # 利用网络的输出outputs和标签labels计算损失值

        loss = criterion(predicts, labels)
        loss.backward()
        optimizer.step()
        # 打印轮次和损失值

        running_loss += loss.item()
        if (i + 1) % 600 == 0:
            print('[%d, %5d] loss: %.5f' %
                  (epoch + 1, i + 1, running_loss / 600))
            running_loss = 0.0
print('Finished Training')

# 查看训练集上的ACC

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device),data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

# 查看在每个分类上的ACC

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device),data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

for i in range(10):
    print('Accuracy of %5s : %2d %%' % (
        classes[i], 100 * class_correct[i] / class_total[i]))

# 展示预测图片和结果

dataiter = iter(testloader)
images, labels = dataiter.next()
# 打印原始图片

imshow(torchvision.utils.make_grid(images))
# 打印真实的标签

print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(32)))

# 模型保存

# 首先设定模型的保存路径

PATH = './cifar_net.pth'
# 保存模型的状态字典

torch.save(model.state_dict(), PATH)

# 模型加载

PATH = './cifar_net.pth'
model.load_state_dict(torch.load(PATH))