MachineLP的Github(欢迎follow):https://github.com/MachineLP
本小节使用torch搭建逻辑回归模型,训练和测试:
(1)定义模型超参数:输入大小、类别数量、迭代次数、批量大小、学习率。
(2)获取训练数据:使用已有的,或者使用自己的数据集:(可参考:https://blog.csdn.net/u014365862/article/details/80506147)
(3)定义模型。
(4)定义损失函数。
(5)定义优化算法。
(6)训练->测试->保存模型。
---------------------------------我是可爱的分割线---------------------------------
代码部分:
# coding=utf-8 import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms # 训练模型的超参数 input_size = 784 num_classes = 10 num_epochs = 5 batch_size = 100 learning_rate = 0.001 # MNIST 数据集 (样本,标签) train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True) test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor()) # 构建数据管道, 使用自己的数据集请参考:https://blog.csdn.net/u014365862/article/details/80506147 train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) # 搭建逻辑回归模型, 主要体现在后边的loss上,使用交叉熵 model = nn.Linear(input_size, num_classes) # 损失函数和优化算法 # nn.CrossEntropyLoss() computes softmax internally criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) # train_loader的大小,也就是含有多少个bach。 total_step = len(train_loader) # 训练模型 # 在整个数据集上迭代的次数 for epoch in range(num_epochs): # 每次取一个bach进行训练。 for i, (images, labels) in enumerate(train_loader): # 将数据reshape到模型需要的大小。 images = images.reshape(-1, 28*28) # 前向传播 outputs = model(images) # 计算模型的loss loss = criterion(outputs, labels) # 后向传播,更新模型参数 optimizer.zero_grad() loss.backward() optimizer.step() # 每100个batch,打印一次结果 if (i+1) % 100 == 0: print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(epoch+1, num_epochs, i+1, total_step, loss.item())) # 模型测试 # 在测试阶段,不用计算梯度 with torch.no_grad(): correct = 0 total = 0 for images, labels in test_loader: images = images.reshape(-1, 28*28) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum() print('Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total)) # 保存模型 torch.save(model.state_dict(), 'model.ckpt')
总结:
torch搭建模型训练的流程和TF基本一致,相比来说torch更简洁一些,TF稍微繁琐一些,单这不影响你是TF的铁粉。
如何你想训练其他模型,只需要把上面模型部分替换掉就可以了,这是一个简单的实现,里边有很多功能没加,想更快的提升自己的能力,可以考虑自己基于torch搭建一个通用训练框架。
torch系列: