pytorch安装

python3.6

pip3 install torchvision==0.2.0

MNIST手写数字0-9训练集下载
链接：https://pan.baidu.com/s/1kUiARaVTdFNAfJoS5xx6pw
提取码：xttr

错误

raise NotSupportedError(base.range(), "slicing multiple dimensions at the sa

借用一下其他大佬的测试代码

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms


# 设置网络结构
class Net(nn.Module):
    def __init__(self, in_features, out_features):
        super(Net, self).__init__()
        self.dnn1 = nn.Linear(in_features, 512)  # 第一层全连接层
        self.dnn2 = nn.Linear(512, out_features)  # 第二层全连接层

    def forward(self, x):
        x = F.relu(self.dnn1(x))  # relu激活
        x = self.dnn2(x)
        return x


net = Net(28 * 28, 10)

# 加载MNIST数据集
train_dataset = datasets.MNIST('../data/MNIST', train=True, download=False, transform=transforms.ToTensor())
test_dataset = datasets.MNIST('../data/MNIST', train=False, download=False, transform=transforms.ToTensor())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)
# print(net)

# 开始训练
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.02)

for epoch in range(5):
    running_loss, running_acc = 0.0, 0.0
    for i, data in enumerate(train_loader):
        img, label = data
        img = img.reshape(-1, 28 * 28)
        out = net(img)
        loss = criterion(out, label)

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

        running_loss += loss.item() * label.size(0)
        _, predicted = torch.max(out, 1)
        running_acc += (predicted == label).sum().item()
        print('Epoch [{}/5], Step [{}/{}], Loss: {:.6f}, Acc: {:.6f}'.format(
            epoch + 1, i + 1, len(train_loader), loss.item(), (predicted == label).sum().item() / 128))
    # 测试
    test_loss, test_acc = 0.0, 0.0
    for i, data in enumerate(test_loader):
        img, label = data
        img = img.reshape(-1, 28 * 28)
        out = net(img)
        loss = criterion(out, label)

        test_loss += loss.item() * label.size(0)
        _, predicted = torch.max(out, 1)
        test_acc += (predicted == label).sum().item()

    print("Train {} epoch, Loss: {:.6f}, Acc: {:.6f}, Test_Loss: {:.6f}, Test_Acc: {:.6f}".format(
        epoch + 1, running_loss / (len(train_dataset)), running_acc / (len(train_dataset)),
        test_loss / (len(test_dataset)), test_acc / (len(test_dataset))))

效果