测试结果最后两行分别为预测类别与真实类别。数据预览这里的数据使用的是mnist数据集,大家可以将代码中的DOWNLOAD_MNIST值修改为True进行自动下载。代码import torchimport torch.nn as nnimport torch.utils.data as Dataimport torchvision # 数据库模块import matplotlib.pyplot as plt#训练整批数据多少次,这里为了节约时间,只训练一次EPOCH=1#每次批处理50个数据BATCH_SIZE=50#学习效率LR=0.001# 如果已经下载好了mnist数据就写上FalseDOWNLOAD_MNIST = False #训练的数据集:Mnist手写数字train_data=torchvision.datasets.MNIST( #保存或提取数据集的位置 root=’./mnist/’, #该数据是训练数据 train=True, #转换PIL.Image or numpy.ndarray成torch.FloatTensor (C x H x W), 训练的时候 normalize 成 [0.0, 1.0] 区间 transform=torchvision.transforms.ToTensor(), #没下载就下载,下载了就不用再下了 download=DOWNLOAD_MNIST,)#绘制一下数据集#黑色的地方的值都是0, 白色的地方值大于0.print(train_data.train_data.size()) # (60000, 28, 28)print(train_data.train_labels.size()) # (60000)plt.imshow(train_data.train_data[2].numpy(), cmap=‘gray’)plt.title(’%i’ % train_data.train_labels[2])plt.show()#测试数据test_data=torchvision.datasets.MNIST(root=’./mnist/’,train=False)#批训练50samples,1 channel,28x28 (50, 1, 28, 28)train_loader=Data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True)#这里只测试了前2000个#特征test_x=torch.unsqueeze(test_data.test_data,dim=1).type(torch.FloatTensor)[:2000]/255.#标签test_y=test_data.test_labels[:2000]#构建CNN模型class CNN(nn.Module): def init(self): super(CNN,self).init() #input shape(1,28,28) self.conv1=nn.Sequential( #卷积 nn.Conv2d( in_channels=1, out_channels=16, #filter size kernel_size=5, #filter movement/step stride=1, #如果想要con2d出来的图片长宽没有变化, #padding=(kernel_size-1)/2当stride=1 padding=2, ), #output shape(16,28,28) #激励函数 nn.ReLU(), #池化 # 在2x2空间里向下采样,output shape(16,14,14) nn.MaxPool2d(kernel_size=2), ) #input shape(16,14,14) self.conv2=nn.Sequential( nn.Conv2d(16,32,5,1,2), #output shape(32,14,14) #激励函数 nn.ReLU(), #output shape(32,7,7) nn.MaxPool2d(2), ) #全连接层——进行分类。这里将其分成了10类 self.out=nn.Linear(3277,10) def forward(self,x): x=self.conv1(x) x=self.conv2(x) #展平多维的卷积图成(batch_size,3277) x=x.view(x.size(0),-1) output=self.out(x) return outputcnn=CNN()print(cnn)#训练#优化器optimizer=torch.optim.Adam(cnn.parameters(),lr=LR)#损失函数loss_func=nn.CrossEntropyLoss()#开始训练for epoch in range(EPOCH): for step,(b_x,b_y) in enumerate(train_loader): #将数据输入nn并且得到output output=cnn(b_x) #计算output与真实值之间的误差 loss=loss_func(output,b_y) #清空上一步残余更新参数值 optimizer.zero_grad() #误差反向传播,让参数进行更新 loss.backward() #将更新后的参数值施加到nn的parameters上 optimizer.step()#测试:选取10个数据test_output=cnn(test_x[:10])pred_y=torch.max(test_output,1)[1].data.numpy().squeeze()print(pred_y, ‘prediction number’)print(test_y[:10].numpy(), ‘real number’)# if name==’main’:# print(“hello word”)