#导入所用库
import os
import torch
import torchvision
from torch import optim, nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import torch.nn.functional as F
#获得数据
def get_data():
transf = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5],0.5)]) #定义对数据的处理，totensor将数据转换成pytorch计算的tensor，normalize对数据进行归一化
train_dataset = datasets.MNIST(root='./data',train= True,transform=transf,download= False)#定义train_dataset，download，在第一次运行时候设置为True。
test_dataset = datasets.MNIST(root='./data', train=False, transform=transf, download=False)#定义test_dataset，download，在第一次运行时候设置为True。
train_dataloader = DataLoader(train_dataset,batch_size=64,shuffle=True,drop_last=True)#定义获取数据方式：dataloader
test_dataloder = DataLoader(test_dataset,batch_size = 1,shuffle = False,drop_last=True)
return train_dataloader,test_dataloder
#定义网络模型
#网络模型为三层卷积两层全连接网络，每层卷积之后跟着最大池化和激活
class CONV_NET(nn.Module): #定义卷积模型类，需要继承nn.Module
def __init__(self): #构造函数
super(CONV_NET, self).__init__() #先调用父类构造函数
self.conv1 = nn.Conv2d(1, 10, kernel_size = 5) # 定义第一类卷积层，输入通道为1，输出通道为10，卷积核大小为5
self.conv2 = nn.Conv2d(10, 20, kernel_size=3) # 定义第二类卷积层，输入通道为10，输出通道为20，卷积核大小为3
self.conv3 = nn.Conv2d(20, 50, kernel_size=2) # 定义第三类卷积层，输入通道为20，输出通道为50，卷积核大小为3
self.conv2_drop = nn.Dropout2d() # 二维dropout
self.conv3_drop = nn.Dropout2d()
self.fc1 = nn.Linear(200, 50) # 第一个全连接层 输入维度200，输出50
self.fc2 = nn.Linear(50, 10) # 第二个全连接层 输入维度50，输出10
def forward(self, x): #定义调用模型时数据流向，x为输入的图像
x = F.relu(F.max_pool2d(self.conv1(x), 2)) #第一个卷积，池化，激活 #输入batch 1 28 28，卷积输出batch 10 24 24 池化输出batch 10 12 12
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))#第二个卷积，池化，激活#输入batch 10 12 12，卷积输出 输出batch 20 10 10，池化输出batch 20 5 5
# drop不改变维度随机删去一些数据
x = F.relu(F.max_pool2d(self.conv3_drop(self.conv3(x)), 2))#第三个卷积，池化，激活 #输入 batch 20 5 5 卷积输出 batch 50 4 4 池化 batch 50 2 2
x = x.view(-1, 200) # view将数据变为200维的数据 batch 50 2 2变为batch 200
x = F.relu(self.fc1(x)) # 第一层全连接层 输入维度 batch 200 输出batch 50
x = F.dropout(x, training=self.training) #dropout层，防止过拟合
x = self.fc2(x) # 第二层全连接层 输入维度 batch 50 输出batch 10
return x # 输出
if __name__ =="__main__": #main 函数，程序入口
# 指定超参数
batch_size = 64 #每次处理batch 大小
lr = 1e-2 #学习率
epochs = 3 #训练轮数
log_dir = './logs/' #日志目录
writer = SummaryWriter(log_dir) #定义日志记录对象
momentum = 0.5 # 随机梯度下降优化的动量系数
train_loader , test_loader = get_data() #获取训练数据和测试数据
model = CONV_NET() #构建模型对象
if os.path.isfile('./model.pth'): #判断当前目录下model.pth是否存在，model.pth为本次实验保存的模型
model_dict = torch.load('./model.pth') #加载模型，读取为字典
model.load_state_dict(model_dict) #将字典的数据赋值为网络模型的参数
print(model) #打印模型
print('load success') #输出提示
loss_func = nn.CrossEntropyLoss() #定义损失函数
optimizer = optim.SGD(model.parameters(),lr= lr,momentum=momentum) #定义优化器
model.train() #将模型设置为训练模式
batch = 0 #batch为计数变量
for i in range(epochs): #循环多少轮
for data in train_loader: #利用train_dataloader拿数据
img ,label = data # 拿到的数据分别是img和label
image = torchvision.utils.make_grid(img) #将一个batch的img展开成一张图像 为了后面写入tensorboard
out = model(img) #将img输入到模型中得到输出
loss = loss_func(out,label) #计算输出和真实值之间的误差
optimizer.zero_grad() #将当前网络的梯度置为o
loss.backward() #误差反向传播，计算网络每层梯度
optimizer.step() #按照梯度方向，更新网络参数
batch += 1 #计数变量加1
print('epoch = {},batch = {}, loss = {}'.format(i,batch, loss)) #打印当前轮数，训练次数，损失
writer.add_scalar('loss_value',loss,global_step=batch) #使用tensorboard 记录损失
# writer.add_image('batch_image',image,global_step=batch) # 使用tensorboard 记录当前参与计算的图像
torch.save(model.state_dict(),'./model.pth') #每训练一轮，保存一次模型
writer.add_graph(model, input_to_model=img, ) #tensorboard 记录当前模型计算图
writer.close() #关闭tensorboard记录
model.eval() #模型设置为写模式
count = 0 #记录分类正确的个数
for img ,label in test_loader: #每次拿出一个图像进行测试
out = model(img) #单个图像得到分类置信度
_,predict = torch.max(out,1) #根据置信度使用max函数将置信度最大的类别作为分类结果
if predict == label: #判断预测分类结果和真实的是否一致
count += 1 #如果一致则分类正确的计数加1
print("acc = {}".format(count/len(test_loader))) #所有测试集数据测试完成后计算准确率
网友评论