import torch.nn as nn
import torch.nn.functional as F
import torch


class LeNet(nn.Module):         # 首先定义一个类 这个类要继承于nn.Module
    def __init__(self):  # 定义初始化函数       # 在这个类中实现两个方法 一 初始化函数 在搭建网络过程中所需要使用到的网络层结构
        super(LeNet, self).__init__()  # super函数 在定义这个类的过程中 继承了nn.module super函数解决在多重继承中调用父类可能会出现的问题
        self.conv1 = nn.Conv2d(3, 16, 5) # （in_channels输入通道数, out_channels也是卷积核数量， kernel_size卷积核尺寸                                      # 定义第一个卷积层 就是用nn.Conv2d 内部定义方法
        self.pool1 = nn.MaxPool2d(2, 2) # (kernel_size池化核大小，stride or kernel_size 如果不指定就是池化核大小)
        self.conv2 = nn.Conv2d(16, 32, 5) #
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):   # 定义正向传播过程  实例化这个类之后 将参数传递到这个实例中 就会进行正向传播
        x = F.relu(self.conv1(x))   # input(3 x 32 x 32)   #output(16 x(32 - 5 + 0x2)/1+1  x 28 )
        x = self.pool1(x)             #   output (16 x 14 x 14)
        x = F.relu(self.conv2(x))   # input(16 x 14 x 14)  output(32 x 10 x 10)
        x = self.pool2(x)        # output(32 x 5 x 5)
        x = x.view(-1, 32 * 5 * 5)  # 展成一维向量 -1为自动填充
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


input1 = torch.rand([32, 3, 32, 32])
model = LeNet()
print(model)
output = model(input1)

model

train

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import  torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# transform.Compose 将预处理方法打包成一个整体  1.ToTensor 将pil图像或者np数据转化成tensor 2.normalize 使用均值和标准差 来标准化tensor
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)

# root='./' 把数据集下载到什么地方 一般放在当前目录的data文件夹下 train=True 会导入数据集中训练集  transform为图像预处理
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=36, shuffle=True, num_workers=0)
# 将刚才的数据集导入 batch_size 每一批随机拿出36张拿来训练  shuffle打乱 num_workers=0 不用多线程


# 导入10000张测试图片
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=10000, shuffle=False, num_workers=0)

test_data_iter = iter(testloader)  # iter 将testloader转化成可迭代的迭代器
test_image, test_label = test_data_iter.next()  # 通过next方法 获取一批数据包括测试的图像 以及图像的标签值  这个非常好用！

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 元组类型 值是不能带变的 对应的标签 index0 就是plane



# def imshow(img):
#     img = img / 2 + 0.5   # unnormalize  刚才标准化现在还原
#     npimg = img.numpy()   # 转化成numpy格式
#     plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 刚才是channel height width 现在要变为 height width channel 所以从（0 1 2） 改成（1 2 0）
#     plt.show()  # 用show展示出来
#
# print(' '.join('%5s' % classes[test_label[j]] for j in range(4)))  # print labels
# imshow(torchvision.utils.make_grid(test_image))    # show images

net = LeNet()
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)  # net.parameters()所需要训练的参数

for epoch in range(5):  # 训练集需要训练多少次
    running_loss = 0.0  # 累加loss
    for step, data in enumerate(trainloader, start=0):  # 遍历训练集样本 enumerate 不仅能返回每一批的数据data 还会返回这一批data所对应的步数 也就是index
        # gat the inputs， data is a list of [inputs, labels]
        inputs, labels = data    # 将inputs（输入的图像） labels（标签） 分离出来
        optimizer.zero_grad()   # 将历史的梯度清零 可以一次性计算多个小的batch
        outputs = net(inputs)
        loss = loss_function(outputs, labels)  # 计算损失
        loss.backward()   # 反向传播
        optimizer.step()  # 参数更新 也是权重更新

        running_loss += loss.item()  # 每次计算完loss 累加到running_loss变量中
        if step % 500 == 499:      # print every 500 mini_batch
            with torch.no_grad():    # torch.no_grad 在接下来过程中 不要计算每个节点的误差损失梯度 少占用内存或者资源
                outputs = net(test_image)  # 正向传播     # 这里的[1] 代表我们只需要他的index（索引） 只需要知道在哪个地方
                predict_y = torch.max(outputs, dim=1)[1]  #toech.max 寻找输出的最大的index在哪个位置 网络预测最可能属于哪个类别
                # dim = 1 在维度1 上寻找最大的值 [batch, 10]   因为第0个维度对应的是batch 要在输出的十个节点中寻找
                accuracy = (predict_y == test_label).sum().item() / test_label.size(0)   # 计算出来的是个tensor 要通过item得出数值
                # 将预测的标签类别和真实的标签类别进行比较 再求和 在本次测试中预测对了多少个样本
                print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
                      (epoch + 1, step + 1, running_loss / 500, accuracy))  # step 在每一轮中的多少步
                running_loss = 0.0  # 将running_loss 清零 进行下一个500步的训练 并计算loss

print('Finished Training')

save_path = './Lenet.pth'  # 将模型保存
torch.save(net.state_dict(), save_path)

predict

import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

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

net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))  # load_state_dict载入保存的权重文件

im = Image.open('5.jpg')  # pillow 载入图像
im = transform(im)   # C H W 转化成N C H W
im = torch.unsqueeze(im, dim=0)  # unsqueeze 在 0 前面再加上一个维度

with torch.no_grad():
    outputs = net(im)
    predict = torch.max(outputs, dim=1)[1].data.numpy()   # 将index传入classes
print(classes[int(predict)])

完结