逻辑回归的pytorch实现

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

#设置相关参数
input_size=28*28
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())

#数据加载
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)

#定义model、损失函数与优化函数
model=nn.Linear(input_size,num_classes)
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(model.parameters(),lr=learning_rate)

#开始迭代训练
total_step=len(train_loader)
for epoch in range(num_epochs):
    for i,(images,labels) in enumerate(train_loader):
        images=images.reshape(-1,input_size)
        outputs=model(images)
        loss=criterion(outputs,labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        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,input_size)
        outputs=model(images)
        _,predicted=torch.max(outputs.data,1)
        total=total+labels.size(0)
        correct=correct+(predicted==labels).sum()

    print("Accuracy of the model on the 10000 test images:{} %".format(100*correct/total))

#保存模型
torch.save(model.state_dict(),"logisticModel.ckpt")

运行的时候会先下载相关数据集，随后开始训练建模和测试模型的准确率。最终结果如下：

Epoch [4/5], Step [100/600], Loss:1.2736
Epoch [4/5], Step [200/600], Loss:1.1791
Epoch [4/5], Step [300/600], Loss:1.2005
Epoch [4/5], Step [400/600], Loss:1.2209
Epoch [4/5], Step [500/600], Loss:1.0825
Epoch [4/5], Step [600/600], Loss:1.0868
Epoch [5/5], Step [100/600], Loss:1.2321
Epoch [5/5], Step [200/600], Loss:0.9939
Epoch [5/5], Step [300/600], Loss:1.0931
Epoch [5/5], Step [400/600], Loss:1.1036
Epoch [5/5], Step [500/600], Loss:1.0602
Epoch [5/5], Step [600/600], Loss:0.9724
Accuracy of the model on the 10000 test images:82 %

 推荐一个公众号：健哥聊量化，会持续推出股票相关基础知识，以及python实现的一些基本的分析代码。欢迎大家关注，二维码如下：