import torch
import torch.utils.data as Data
from torchvision import datasets,transforms
import torchvision
from torch.autograd import  Variable
import numpy as np
import matplotlib.pyplot as plt

torchvision的包的主要功能是实现数据的处理，导入和预览呢，所以如果需要对计算机视觉的相关问题 进行处理，就可以利用 torchvision包中提供的大量的类来完成相应的工作。
torchvision.datasets加需要下载的数据集名称就可以了，而且torchvision还包含了COCO、imageNet、CIFCAR等数据集。

transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,],std=[0.5,])])

数据标准化 ToTensor可以变为0-1的，后面一句变为-1到1，RGB每一个通道都执行（x-0.5）/0.5

data_train=datasets.MNIST(root="D:\jupyter_data",  transform=transform, train=True,
                          download=True
                          )
data_test=datasets.MNIST(root="D:\jupyter_data", transform=transform, train=False)

1.按批获取数据

batch = 256
train_iter = Data.DataLoader(dataset=data_train, batch_size=batch, shuffle=True)
test_iter = Data.DataLoader(data_test,batch, shuffle=True)
print(train_iter)

<torch.utils.data.dataloader.DataLoader object at 0x0000021B52AF2B00>

2.定义和初始化模型

for X,y in train_iter:
    print(X.shape)
    print(X.view(X.shape[0],-1).shape)
    break

torch.Size([256, 1, 28, 28])
torch.Size([256, 784])

num_inputs = 784
num_outputs = 10
class LinearNet(torch.nn.Module):
    def __init__(self,num_inputs,num_outputs):
        super(LinearNet,self).__init__()
        self.linear = torch.nn.Linear(num_inputs,num_outputs)
    def forword(self,X): #这里X是28行28列.
        y=self.linear(X.view(X.shape[0],-1))
        return y  #参数中的-1就代表这个位置由其他位置的数字来推断,根据上面得到的shape，一批是256个，所以这里的shape[0]是256
net = LinearNet(num_inputs, num_outputs)
torch.nn.init.normal_(net.linear.weight, mean=0, std=0.01)
torch.nn.init.constant_(net.linear.bias, val=0)

Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)

• 看一下输出的结构，因为用softmax函数，所以要符合softmax的输入

print(net)
for X,y in train_iter:
    y_pre = net.forword(X)
    print(y_pre.shape)
    print(y_pre.view(y_pre.shape[0],-1).shape)
    break

LinearNet(
  (linear): Linear(in_features=784, out_features=10, bias=True)
)
torch.Size([256, 10])
torch.Size([256, 10])

3.定义损失函数。这里用softmax和交叉熵合起来的

loss = torch.nn.CrossEntropyLoss()

4.定义优化算法

optimizer = torch.optim.SGD(net.linear.parameters(),lr=0.1)

5.训练模型

epoch_size = 3
def softmax_train(train_iter, test_iter, epoch_size, loss, optimizer, net, batchsize):
    for epoch in range(0, epoch_size+1):
        train_acc_count = 0
        train_loss = 0
        n = 0
        for X,y in train_iter:
            y_pre = net.forword(X)
            l = loss(y_pre, y).sum() #注意要加起来，这个是softmax+交叉熵的，之前的线性回归不用
            optimizer.zero_grad() #梯度清零
            l.backward()
            optimizer.step()
            #准确率
            train_acc_count += (y_pre.argmax(dim=1) == y).sum().item()
            train_loss += l.item() #
            n += y.shape[0]
        print('epoch %d,loss %.4f, train acc %.3f'% (epoch + 1, train_loss/n ,train_acc_count/n))
softmax_train(train_iter, test_iter, epoch_size, loss, optimizer, net, batch)

epoch 1,loss 0.0025, train acc 0.824
epoch 2,loss 0.0014, train acc 0.890
epoch 3,loss 0.0013, train acc 0.899
epoch 4,loss 0.0013, train acc 0.904

6.查看测试集的准确率

def test_acc(test_iter, net):
    test_acc_count = 0
    n = 0
    for X,y in test_iter:
        y_pre = net.forword(X)
        test_acc_count += (y_pre.argmax(dim=1) == y).sum().item()
        n += y.shape[0]
    print('test acc %.3f'% (test_acc_count/n))
test_acc(test_iter, net)

test acc 0.910