Python MLP与CNN 对MNIST数据集分类

前言MNIST数据集介绍导入MNIST数据集可视化数据展示图像细节定义MLP神经网络结构定义CNN神经网络结构定义损失函数和梯度下降优化器训练MLP神经网络训练CNN神经网络可视化训练误差和验证误差载入验证集损失最低的模型整个测试集验证MLP可视化数据CNN可视化数据

前言

  本文主要是学习记录,在小土堆Pytorch课程学习后第一次写,对于MLP与CNN的相关知识还不够了解,后续对这块学习时再进行补充。这里特别感谢 小土堆Pytorch课程、同济子豪兄的MNIST课程。

MNIST数据集介绍

  MNIST 的全称是 Modified National Institute of Standards and Technology database。   MNIST数据集是一个常用的手写数字识别数据集,包含了大约 60000 张训练集图片和 10000 张测试集图片,每张图片都是 28 像素 * 28 像素的灰度图像。这些图像都经过了预处理和标准化,使得每张图像都被表示为一个行向量,其中每个元素的值都在 0 到 1 之间。MNIST数据集中的每个图像都标注有对应的数字,因此该数据集通常用于训练和评估机器学习算法和模型的性能,尤其是对手写数字识别算法和模型的评估。

导入MNIST数据集

Pytorch官方介绍:Pytorch官网MNIST文档

import torchvision.transforms as transforms

from torchvision import datasets

from torch.utils.data.sampler import SubsetRandomSampler

#在数据中使用的处理器个数

num_workers = 0

#每一批数据的个数

batch_size = 20

#验证数据集的占比

valid_size = 0.2

#将数据转化为Pytorch的张量Tensor类型

transform = transforms.ToTensor()

# 选择训练集和测试集,root:数据文件根目录

# train:是否载入训练集

# download:如果目录中找不到数据集,是否自动下载

# transform:将载入的数据按上面transform定义的方式进行转换

##训练数据下载

train_data = datasets.MNIST(root="./MNIST_Dataset",train=True,

transform=transform,download=True)

##测试数据下载

test_data = datasets.MNIST(root="./MNIST_Dataset",train=False,

transform=transform,download=True)

这里将加载的训练数据集进行划分,分为训练和验证这样方便对最后测试集进行最后的检验 测试集相当于是最后考试,划分为的验证集相当于是模拟测试,训练集则为学习

#### 将训练集打乱并分为验证集索引和训练集索引 ###

num_train = len(train_data) # 获取训练集的长度

indices = list(range(num_train)) # 生成list,里面从0开始一共有num_train长度

np.random.shuffle(indices) # 将列表打乱(相当于随即将训练集打乱,类似于DataLoader的shuffle)

split = int(np.floor(valid_size*num_train)) # 计算分化split长度

train_idx = indices[split:None] # 获取训练集索引(list)

valid_idx = indices[None:split] # 获取验证集索引(list)

#### 将训练数据集划分为新的训练数据集和验证集 ###

train_sampler = SubsetRandomSampler(train_idx) # 相当于将list转换为Sampler,并随机采样序列

valid_sampler = SubsetRandomSampler(valid_idx) # 相当于将list转换为Sampler,并随机采样序列

### 创建数据Loader,python中的生成器,每次条用返回一个batch ###

# 训练集

train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,

sampler=train_sampler, num_workers=num_workers)

# 验证集

valid_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,

sampler=valid_sampler, num_workers=num_workers)

# 测试集

test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size

,shuffle= True, num_workers=num_workers)

可视化数据

# 可重复运行,展示不同的图片

# 构造迭代器,获得训练集中的一批次的数据

dataiter = iter(train_loader)

images, labels = next(dataiter)

# 将数据集中的Tensor张量转换为numpy的array数据类型

images = images.numpy()

# 可视化图片和标签

fig = plt.figure(figsize=(15,3))

for idx in np.arange(20):

ax = fig.add_subplot(2, 10, idx+1, xticks=[], yticks=[])

ax.imshow(np.squeeze(images[idx]),cmap='gray')

# .item()获张量的数值

ax.set_title(str(labels[idx].item()))

展示图像细节

# 去掉图像的批次维度,只保留索引为1的单张图像的长宽像素值

img = np.squeeze(images[1])

fig = plt.figure(figsize=(12,12))

ax = fig.add_subplot(111)

ax.imshow(img, cmap='gray')

width, height = img.shape

thresh = img.max()/2.5

# 遍历每一行每一列每一个元素

for x in range(width):

for y in range(height):

# 像素保留两位小数,如果为0则显示为0

val = round(img[x][y],2) if img[x][y] != 0 else 0

# annotate 在每个像素上标记,在0上为白色,在亮的地方为黑色

ax.annotate(str(val), xy=(y,x),

horizontalalignment='center',

verticalalignment='center',

color='white' if img[x][y]

定义MLP神经网络结构

输入为784维向量,输出为10个数字对应的概率 ,中间层是两个隐含层,每个层都是512

import torch.nn as nn

class Net(nn.Module):

def __init__(self):

super(Net,self).__init__()

self.model = nn.Sequential(

nn.Linear(28*28, 512), # 输入层到隐含层 784 -> 512

nn.ReLU(), # 非线性激活ReLU

nn.Dropout(0.3), # Dropout 防止过拟合

nn.Linear(512, 512), # 隐含层到隐含层 512 -> 512

nn.ReLU(), # 非线性激活ReLU

nn.Dropout(0.3), # Dropout 防止过拟合

nn.Linear(512, 10), # 隐含层到输出层 512 ->10

)

def forward(self,x):

flatten = nn.Flatten()

x = flatten(x)

x = self.model(x)

return x

Model = Net()

if torch.cuda.is_available():

Model = Model.cuda()

print(Model)

Net(

(model): Sequential(

(0): Linear(in_features=784, out_features=512, bias=True)

(1): ReLU()

(2): Dropout(p=0.3, inplace=False)

(3): Linear(in_features=512, out_features=512, bias=True)

(4): ReLU()

(5): Dropout(p=0.3, inplace=False)

(6): Linear(in_features=512, out_features=10, bias=True)

)

)

定义CNN神经网络结构

如图所示,大致卷积大致结构为:

import torch.nn as nn

class CNNet(nn.Module):

def __init__(self):

super().__init__()

self.model = nn.Sequential(

nn.Conv2d(in_channels= 1 ,out_channels= 5 ,kernel_size= 5,padding=2), # 卷基层1

nn.ReLU(), # 激活函数

nn.MaxPool2d(kernel_size=2), # 最大池化

nn.Conv2d(in_channels=5 ,out_channels= 16,kernel_size= 5, padding=2), # 卷基层2

nn.ReLU(), # 激活函数

nn.MaxPool2d(kernel_size=2), # 最大池化

nn.Flatten(), # 展开

nn.Linear(784,64),

nn.ReLU(),

nn.Linear(64,10)

)

def forward(self,x):

x = self.model(x)

return x

Model_CNN = CNNet()

if torch.cuda.is_available():

Model_CNN = Model_CNN.cuda()

print(Model_CNN)

CNNet(

(model): Sequential(

(0): Conv2d(1, 5, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))

(1): ReLU()

(2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)

(3): Conv2d(5, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))

(4): ReLU()

(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)

(6): Flatten(start_dim=1, end_dim=-1)

(7): Linear(in_features=784, out_features=64, bias=True)

(8): ReLU()

(9): Linear(in_features=64, out_features=10, bias=True)

)

)

定义损失函数和梯度下降优化器

使用交叉熵作为分类问题的损失函数

CrossEntropyLoss -> Pytorch交叉熵文档Optim -> Pytorch优化器文档

# 损失函数为交叉熵损失函数

criterion = nn.CrossEntropyLoss()

if torch.cuda.is_available():

criterion = criterion.cuda()

# 定义优化器

optimizer = torch.optim.SGD(Model.parameters(),lr=0.01)

optimizer_cnn = torch.optim.SGD(Model_CNN.parameters(),lr=0.01)

训练MLP神经网络

1.清除所有梯度 2.正向预测,求出模型对数据集的预测分类 3.计算损失函数 4.反向传播,计算损失函数的每一个权重求导,求得对应权重的梯度 5.优化器更新权重 6.计算每一轮的平均训练误差和验证误差

# 训练轮次,每一轮都完整遍历数据集中所有图像

num_epoch = 20

# 初始化验证集最小误差为无穷大

valid_loss_min = np.Inf

# 初始化loss_list

train_loss_list = []

valid_loss_list = []

# 每一轮训练:

for epoch in range(num_epoch):

#初始化损失

train_loss = 0

valid_loss = 0

### 训练阶段 ###

Model.train()

for data in train_loader:

imgs, labels = data # 读取数据

if torch.cuda.is_available():

imgs = imgs.cuda()

labels = labels.cuda()

optimizer.zero_grad() # 将梯度归零

output = Model(imgs) # 正向预测

loss = criterion(output,labels) # 计算损失

loss.backward() # 反向传播

optimizer.step() # 优化器进行更新

train_loss += loss.item()*imgs.size(0) # 计算本批次损失

### 验证阶段 ###

Model.eval()

for data in valid_loader:

imgs, labels = data # 读取数据

if torch.cuda.is_available():

imgs = imgs.cuda()

labels = labels.cuda()

output = Model(imgs) # 正向预测

loss = criterion(output,labels) # 计算损失

valid_loss += loss.item()*imgs.size(0) # 计算本批次损失

# 结束本论的训练和验证,打印训练和验证的指标

# 计算平均训练损失和平均验证损失,存储在列表中

train_loss = train_loss/len(train_loader.dataset)

valid_loss = valid_loss/len(valid_loader.dataset)

train_loss_list.append(train_loss)

valid_loss_list.append(valid_loss)

print("----------------------------------------------------")

print("第{}轮 \t训练损失: {:.6f} \t验证损失: {:.6f}".format(epoch+1, train_loss, valid_loss))

#损失有降低则保存到本地

if valid_loss <= valid_loss_min:

print("验证误差比上次降低了({:.6f} ---> {:.6f}). 保存模型!".format(valid_loss_min,valid_loss))

torch.save(Model.state_dict(), "MNIST.pth")

valid_loss_min = valid_loss

----------------------------------------------------

第1轮 训练损失: 0.788821 验证损失: 0.077319

验证误差比上次降低了(inf ---> 0.077319). 保存模型!

----------------------------------------------------

第2轮 训练损失: 0.295448 验证损失: 0.059102

验证误差比上次降低了(0.077319 ---> 0.059102). 保存模型!

----------------------------------------------------

第3轮 训练损失: 0.233553 验证损失: 0.048407

验证误差比上次降低了(0.059102 ---> 0.048407). 保存模型!

----------------------------------------------------

第4轮 训练损失: 0.193519 验证损失: 0.040578

验证误差比上次降低了(0.048407 ---> 0.040578). 保存模型!

----------------------------------------------------

第5轮 训练损失: 0.165305 验证损失: 0.035517

验证误差比上次降低了(0.040578 ---> 0.035517). 保存模型!

----------------------------------------------------

第6轮 训练损失: 0.143009 验证损失: 0.031903

验证误差比上次降低了(0.035517 ---> 0.031903). 保存模型!

----------------------------------------------------

第7轮 训练损失: 0.127322 验证损失: 0.028366

验证误差比上次降低了(0.031903 ---> 0.028366). 保存模型!

----------------------------------------------------

第8轮 训练损失: 0.113886 验证损失: 0.026533

验证误差比上次降低了(0.028366 ---> 0.026533). 保存模型!

----------------------------------------------------

第9轮 训练损失: 0.103299 验证损失: 0.024597

验证误差比上次降低了(0.026533 ---> 0.024597). 保存模型!

----------------------------------------------------

第10轮 训练损失: 0.095750 验证损失: 0.022435

验证误差比上次降低了(0.024597 ---> 0.022435). 保存模型!

----------------------------------------------------

第11轮 训练损失: 0.087945 验证损失: 0.021507

验证误差比上次降低了(0.022435 ---> 0.021507). 保存模型!

----------------------------------------------------

第12轮 训练损失: 0.081725 验证损失: 0.019907

验证误差比上次降低了(0.021507 ---> 0.019907). 保存模型!

----------------------------------------------------

第13轮 训练损失: 0.076345 验证损失: 0.019214

验证误差比上次降低了(0.019907 ---> 0.019214). 保存模型!

----------------------------------------------------

第14轮 训练损失: 0.071628 验证损失: 0.018275

验证误差比上次降低了(0.019214 ---> 0.018275). 保存模型!

----------------------------------------------------

第15轮 训练损失: 0.066302 验证损失: 0.017198

验证误差比上次降低了(0.018275 ---> 0.017198). 保存模型!

----------------------------------------------------

第16轮 训练损失: 0.061777 验证损失: 0.016841

验证误差比上次降低了(0.017198 ---> 0.016841). 保存模型!

----------------------------------------------------

第17轮 训练损失: 0.058538 验证损失: 0.015928

验证误差比上次降低了(0.016841 ---> 0.015928). 保存模型!

----------------------------------------------------

第18轮 训练损失: 0.054368 验证损失: 0.016356

----------------------------------------------------

第19轮 训练损失: 0.052580 验证损失: 0.015056

验证误差比上次降低了(0.015928 ---> 0.015056). 保存模型!

----------------------------------------------------

第20轮 训练损失: 0.049286 验证损失: 0.014715

验证误差比上次降低了(0.015056 ---> 0.014715). 保存模型!

训练CNN神经网络

# 训练轮次,每一轮都完整遍历数据集中所有图像

num_epoch_cnn = 20

# 初始化验证集最小误差为无穷大

valid_loss_min_cnn = np.Inf

# 初始化loss_list

train_loss_list_cnn = []

valid_loss_list_cnn = []

# 每一轮训练:

for epoch in range(num_epoch_cnn):

#初始化损失

train_loss = 0

valid_loss = 0

### 训练阶段 ###

Model_CNN.train()

for data in train_loader:

imgs, labels = data # 读取数据

if torch.cuda.is_available():

imgs = imgs.cuda()

labels = labels.cuda()

optimizer_cnn.zero_grad() # 将梯度归零

output = Model_CNN(imgs) # 正向预测

loss = criterion(output,labels) # 计算损失

loss.backward() # 反向传播

optimizer_cnn.step() # 优化器进行更新

train_loss += loss.item()*imgs.size(0) # 计算本批次损失

### 验证阶段 ###

Model_CNN.eval()

for data in valid_loader:

imgs, labels = data # 读取数据

if torch.cuda.is_available():

imgs = imgs.cuda()

labels = labels.cuda()

output = Model_CNN(imgs) # 正向预测

loss = criterion(output,labels) # 计算损失

valid_loss += loss.item()*imgs.size(0) # 计算本批次损失

# 结束本论的训练和验证,打印训练和验证的指标

# 计算平均训练损失和平均验证损失,存储在列表中

train_loss = train_loss/len(train_loader.dataset)

valid_loss = valid_loss/len(valid_loader.dataset)

train_loss_list_cnn.append(train_loss)

valid_loss_list_cnn.append(valid_loss)

print("----------------------------------------------------")

print("第{}轮 \t训练损失: {:.6f} \t验证损失: {:.6f}".format(epoch+1, train_loss, valid_loss))

#损失有降低则保存到本地

if valid_loss <= valid_loss_min_cnn:

print("验证误差比上次降低了({:.6f} ---> {:.6f}). 保存模型!".format(valid_loss_min_cnn,valid_loss))

torch.save(Model_CNN.state_dict(), "MNIST_CNN.pth")

valid_loss_min_cnn = valid_loss

----------------------------------------------------

第1轮 训练损失: 0.526256 验证损失: 0.041250

验证误差比上次降低了(inf ---> 0.041250). 保存模型!

----------------------------------------------------

第2轮 训练损失: 0.110631 验证损失: 0.025260

验证误差比上次降低了(0.041250 ---> 0.025260). 保存模型!

----------------------------------------------------

第3轮 训练损失: 0.077032 验证损失: 0.017001

验证误差比上次降低了(0.025260 ---> 0.017001). 保存模型!

----------------------------------------------------

第4轮 训练损失: 0.061163 验证损失: 0.018689

----------------------------------------------------

第5轮 训练损失: 0.053077 验证损失: 0.013754

验证误差比上次降低了(0.017001 ---> 0.013754). 保存模型!

----------------------------------------------------

第6轮 训练损失: 0.046435 验证损失: 0.013632

验证误差比上次降低了(0.013754 ---> 0.013632). 保存模型!

----------------------------------------------------

第7轮 训练损失: 0.040591 验证损失: 0.014042

----------------------------------------------------

第8轮 训练损失: 0.038105 验证损失: 0.012737

验证误差比上次降低了(0.013632 ---> 0.012737). 保存模型!

----------------------------------------------------

第9轮 训练损失: 0.033946 验证损失: 0.011880

验证误差比上次降低了(0.012737 ---> 0.011880). 保存模型!

----------------------------------------------------

第10轮 训练损失: 0.031287 验证损失: 0.013274

----------------------------------------------------

第11轮 训练损失: 0.028120 验证损失: 0.012257

----------------------------------------------------

第12轮 训练损失: 0.026858 验证损失: 0.010457

验证误差比上次降低了(0.011880 ---> 0.010457). 保存模型!

----------------------------------------------------

第13轮 训练损失: 0.024483 验证损失: 0.015123

----------------------------------------------------

第14轮 训练损失: 0.022900 验证损失: 0.010302

验证误差比上次降低了(0.010457 ---> 0.010302). 保存模型!

----------------------------------------------------

第15轮 训练损失: 0.020943 验证损失: 0.010131

验证误差比上次降低了(0.010302 ---> 0.010131). 保存模型!

----------------------------------------------------

第16轮 训练损失: 0.018959 验证损失: 0.012515

----------------------------------------------------

第17轮 训练损失: 0.018367 验证损失: 0.010460

----------------------------------------------------

第18轮 训练损失: 0.016924 验证损失: 0.009683

验证误差比上次降低了(0.010131 ---> 0.009683). 保存模型!

----------------------------------------------------

第19轮 训练损失: 0.015362 验证损失: 0.011653

----------------------------------------------------

第20轮 训练损失: 0.015022 验证损失: 0.011096

可视化训练误差和验证误差

plt.plot(train_loss_list,label="训练误差")

plt.plot(valid_loss_list,label="验证误差")

plt.legend() # 将label贴上去

plt.title('MLP训练误差和验证误差变化')

plt.show()

plt.plot(train_loss_list_cnn,label="训练误差")

plt.plot(valid_loss_list_cnn,label="验证误差")

plt.legend() # 将label贴上去

plt.title('CNN训练误差和验证误差变化')

plt.show()

plt.plot(valid_loss_list, c='r',label="MLP误差") # c代表就是颜色,r就是指红色

plt.plot(valid_loss_list_cnn, c='g',label="CNN误差") # c代表就是颜色,r就是指红色

plt.legend()

plt.title("CNN/MLP验证误差变化")

plt.show()

载入验证集损失最低的模型

Model.load_state_dict(torch.load('MNIST.pth'))

Model_CNN.load_state_dict(torch.load('MNIST_CNN.pth'))

整个测试集验证

# 遍历整个测试集

# 初始化测试误差

test_loss = 0

class_correct = list(0 for i in range(10))

class_total = list(0 for i in range(10))

# 验证阶段

Model.eval()

for data in test_loader:

imgs, labels = data # 读取数据

if torch.cuda.is_available():

imgs = imgs.cuda()

labels = labels.cuda()

output = Model(imgs) # 正向预测

loss = criterion(output,labels) # 计算损失

test_loss += loss.item()*imgs.size(0) # 计算本批次损失

for i in range(batch_size):

pred = torch.argmax(output[i])

if pred == labels[i].item():

class_correct[labels[i].item()] += 1

class_total[labels[i].item()] +=1

test_loss = test_loss/len(test_loader.dataset)

print('MLP测试集上的误差为:{:.6}'.format(test_loss))

print('MLP整体测试集上正确率为:{}%'.format(((sum(class_correct))/sum(class_total))*100))

MLP测试集上的误差为:0.0689453

MLP整体测试集上正确率为:97.78999999999999%

for i in range(10):

print("MLP: 数字 {} 在测试集上正确率为: {}%".format(i, (class_correct[i]/class_total[i])*100))

MLP: 数字 0 在测试集上正确率为: 98.77551020408163%

MLP: 数字 1 在测试集上正确率为: 99.20704845814979%

MLP: 数字 2 在测试集上正确率为: 97.86821705426357%

MLP: 数字 3 在测试集上正确率为: 98.41584158415841%

MLP: 数字 4 在测试集上正确率为: 97.75967413441956%

MLP: 数字 5 在测试集上正确率为: 96.8609865470852%

MLP: 数字 6 在测试集上正确率为: 97.4947807933194%

MLP: 数字 7 在测试集上正确率为: 97.17898832684824%

MLP: 数字 8 在测试集上正确率为: 97.1252566735113%

MLP: 数字 9 在测试集上正确率为: 96.92765113974232%

# 遍历整个测试集

# 初始化测试误差

test_loss_CNN = 0

class_correct_CNN = list(0 for i in range(10))

class_total_CNN = list(0 for i in range(10))

# 验证阶段

Model.eval()

for data in test_loader:

imgs, labels = data # 读取数据

if torch.cuda.is_available():

imgs = imgs.cuda()

labels = labels.cuda()

output = Model_CNN(imgs) # 正向预测

loss = criterion(output,labels) # 计算损失

test_loss_CNN += loss.item()*imgs.size(0) # 计算本批次损失

for i in range(batch_size):

pred = torch.argmax(output[i])

if pred == labels[i].item():

class_correct_CNN[labels[i].item()] += 1

class_total_CNN[labels[i].item()] +=1

test_loss_CNN = test_loss_CNN/len(test_loader.dataset)

print('CNN测试集上的误差为:{:.6}'.format(test_loss_CNN))

print('CNN整体测试集上正确率为:{}%'.format(((sum(class_correct_CNN))/sum(class_total_CNN))*100))

CNN测试集上的误差为:0.0359796

CNN整体测试集上正确率为:98.83%

for i in range(10):

print("CNN: 数字 {} 在测试集上正确率为: {}%".format(i, (class_correct_CNN[i]/class_total_CNN[i])*100))

CNN: 数字 0 在测试集上正确率为: 99.79591836734694%

CNN: 数字 1 在测试集上正确率为: 99.73568281938327%

CNN: 数字 2 在测试集上正确率为: 98.25581395348837%

CNN: 数字 3 在测试集上正确率为: 98.8118811881188%

CNN: 数字 4 在测试集上正确率为: 98.98167006109979%

CNN: 数字 5 在测试集上正确率为: 98.76681614349776%

CNN: 数字 6 在测试集上正确率为: 98.01670146137788%

CNN: 数字 7 在测试集上正确率为: 99.0272373540856%

CNN: 数字 8 在测试集上正确率为: 98.56262833675564%

CNN: 数字 9 在测试集上正确率为: 98.21605550049554%

MLP可视化数据

# 构造迭代器,获得训练集中的一批次的数据

dataiter = iter(test_loader)

images, labels = next(dataiter)

if torch.cuda.is_available():

images = images.cuda()

labels = labels.cuda()

output = Model(images) # 正向预测

# 将数据集中的Tensor张量转换为numpy的array数据类型

images = images.cpu().numpy()

# 可视化图片和标签

fig = plt.figure(figsize=(15,3))

for idx in np.arange(20):

pred = torch.argmax(output[idx]) #获取预测值

ax = fig.add_subplot(2, 10, idx+1, xticks=[], yticks=[])

ax.imshow(np.squeeze(images[idx]),cmap='gray')

# .item()获张量的数值

ax.set_title("{}[{}]".format(str(labels[idx].item()),pred.item()),

c = "green" if labels[idx].item() == pred.item() else 'r')

CNN可视化数据

# 构造迭代器,获得训练集中的一批次的数据

dataiter = iter(test_loader)

images, labels = next(dataiter)

if torch.cuda.is_available():

images = images.cuda()

labels = labels.cuda()

output = Model_CNN(images) # 正向预测

# 将数据集中的Tensor张量转换为numpy的array数据类型

images = images.cpu().numpy()

# 可视化图片和标签

fig = plt.figure(figsize=(15,3))

for idx in np.arange(20):

pred = torch.argmax(output[idx]) #获取预测值

ax = fig.add_subplot(2, 10, idx+1, xticks=[], yticks=[])

ax.imshow(np.squeeze(images[idx]),cmap='gray')

# .item()获张量的数值

ax.set_title("{}[{}]".format(str(labels[idx].item()),pred.item()),

c = "green" if labels[idx].item() == pred.item() else 'r')

精彩文章

评论可见,请评论后查看内容,谢谢!!!
 您阅读本篇文章共花了: