import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transformsBATCH_SIZE=512     # batch_size即每批训练的样本数量
EPOCHS=20          # 循环次数# 让torch判断是否使用GPU，即device定义为CUDA或CPU
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")train_data = datasets.MNIST('./data',train=True,transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,)) # 标准化] ),download=True)
train_loader = DataLoader(train_data,batch_size=BATCH_SIZE,shuffle=True)test_data = datasets.MNIST('./data',train=False,transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,)) # 标准化] ),download=True)
test_loader = DataLoader(test_data,batch_size=BATCH_SIZE,shuffle=True)#定义神经网络
class ConvNet(nn.Module):def __init__(self):super(ConvNet, self).__init__()#128*128self.conv1 = nn.Conv2d(1,10,kernel_size=5)  #灰度图片的通道1self.conv2 = nn.Conv2d(10,20,kernel_size=3)self.fc1 = nn.Linear(20*10*10,500)self.fc2 = nn.Linear(500,10)def forward(self,x):in_size = x.size(0)  # in_size 为 batch_size（一个batch中的Sample数）x = self.conv1(x)     #输入batch_size*28*28 输出batch_size*24*24x = F.relu(x)x = F.max_pool2d(x,2,2)     #batch_size*14*14x = self.conv2(x)   #batch_size*20*10*10x = F.relu(x)x = x.view(in_size,-1)  #拉伸 20*10*10=2000x = self.fc1(x)     #输入2000  输出batch*500x= F.relu(x)x = self.fc2(x)# softmaxoutput = F.log_softmax(x, dim=1)     #计算分类后每个数字的概率# 返回值 outputreturn output
#定义优化器
model = ConvNet().to(DEVICE)
optimizer = optim.Adam(model.parameters())
#训练
def train(model,device,train_loader,optimizer,epoch):model.train()for batch_idx, (data, target) in enumerate(train_loader):# print(len(train_loader))# print("train_loader_datasets",len(train_loader.dataset))# print("batch_idx:",batch_idx,"data.shape",data.shape,"data_len:",len(data))data, target = data.to(device), target.to(device)  # CPU转GPUoptimizer.zero_grad()  # 优化器清零output = model(data)  # 由model，计算输出值loss = F.cross_entropy(output, target)  # 计算损失函数losspred = output.max(1,keepdim=True)loss.backward()  # loss反向传播optimizer.step()  # 优化器优化if (batch_idx + 1) % 30 == 0:  # 输出结果print("batch_idx:", batch_idx, "batch_idx * len(data)", batch_idx * len(data), "data_len:", len(data),"len(train_loader):",len(train_loader),"train_loader_datasets",len(train_loader.dataset))print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(epoch, batch_idx * len(data), len(train_loader.dataset),100. * batch_idx / len(train_loader), loss.item()))# ---------------------测试函数------------------------------
# 测试的操作也一样封装成一个函数
def test(model, device, test_loader):model.eval()test_loss = 0                           # 损失函数初始化为0correct = 0                             # correct 计数分类正确的数目with torch.no_grad():           # 表示不反向求导（反向求导为训练过程）for data, target in test_loader:    # 遍历所有的data和targetdata, target = data.to(device), target.to(device)   # CPU -> GPUoutput = model(data)            # output为预测值，由model计算出test_loss += F.cross_entropy(output, target).item()     ### 将一批的损失相加pred = output.max(1, keepdim=True)[1]       ### 找到概率最大的下标# pred = torch.argmax(output,dim=1)# print("pred",pred)correct += pred.eq(target.view_as(pred)).sum().item()test_loss /= len(test_loader.dataset)    # 总损失除数据集总数print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(test_loss, correct, len(test_loader.dataset),100. * correct / len(test_loader.dataset)))# 下面开始训练，这里就体现出封装起来的好处了，只要写两行就可以了
# 整个数据集只过一遍
for epoch in range(1,EPOCHS + 1):train(model,DEVICE,train_loader,optimizer,epoch)test(model,DEVICE,test_loader)torch.save(model,"shmodel.pth")