1 导入包

import torchimport numpy as npimport torch.nn as nnfrom torch.utils.data import TensorDataset,DataLoaderimport torchvisionfrom IPython import displayfrom torchvision import transforms

2 加载数据

mnist_train = torchvision.datasets.MNIST(root=\'./datasets/\',download = True,train = True,transform = transforms.ToTensor())mnist_test = torchvision.datasets.MNIST(root=\'./datasets/\',download = True,train = False,transform = transforms.ToTensor())batch_size = 256train_iter = DataLoader(dataset = mnist_train,shuffle = True,batch_size = batch_size)test_iter = DataLoader(dataset = mnist_test,shuffle = True,batch_size = batch_size)

3 定义平铺层

class FlattenLayer(torch.nn.Module):def __init__(self):super(FlattenLayer, self).__init__()def forward(self, x):return x.view(x.shape[0],784)

4 模型选择

num_input,num_hidden1,num_hidden2,num_output = 28*28,512,256,10def choose_model(model_type):if model_type ==\'ReLU\':activation = nn.ReLU()elif model_type ==\'ELU\':activation = nn.ELU()else:activation = nn.Sigmoid()model = nn.Sequential()model.add_module(\"flatten\",FlattenLayer())model.add_module(\"linear1\",nn.Linear(num_input,num_hidden1))model.add_module(\"activation\",activation)model.add_module(\"linear2\",nn.Linear(num_hidden1,num_hidden2))model.add_module(\"activation\",activation)model.add_module(\"linear3\",nn.Linear(num_hidden2,num_output))return modelmodel = choose_model(\'ReLU\')print(model)

5 参数初始化

# for param in model.parameters():#     nn.init.normal_(param,mean=0,std=0.001)for m in model.modules():if isinstance(m, nn.Linear):nn.init.xavier_uniform_(m.weight)nn.init.constant_(m.bias, 0.1)

6 定义训练函数

def train(net,train_iter,test_iter,loss,num_epochs):train_ls,test_ls,train_acc,test_acc = [],[],[],[]for epoch in range(num_epochs):train_ls_sum,train_acc_sum,n = 0,0,0for x,y in train_iter:y_pred = net(x)l = loss(y_pred,y)optimizer.zero_grad()l.backward()optimizer.step()train_ls_sum +=l.item()train_acc_sum += (y_pred.argmax(dim = 1)==y).sum().item()n += y_pred.shape[0]train_ls.append(train_ls_sum)train_acc.append(train_acc_sum/n)test_ls_sum,test_acc_sum,n = 0,0,0for x,y in test_iter:y_pred = net(x)l = loss(y_pred,y)test_ls_sum +=l.item()test_acc_sum += (y_pred.argmax(dim = 1)==y).sum().item()n += y_pred.shape[0]test_ls.append(test_ls_sum)test_acc.append(test_acc_sum/n)print(\'epoch %d, train_loss %.6f,test_loss %f, train_acc %.6f,test_acc %f\'%(epoch+1, train_ls[epoch],test_ls[epoch], train_acc[epoch],test_acc[epoch]))return train_ls,test_ls,train_acc,test_acc_sum

7 定义损失函数和优化器

#训练次数和学习率num_epochs = 20lr = 0.01loss  = nn.CrossEntropyLoss()optimizer = torch.optim.SGD(model.parameters(),lr=lr)

8 模型训练

train_loss,test_loss,train_acc,test_acc = train(model,train_iter,test_iter,loss,num_epochs)