导入必要的包
import torch
from torch import nn
from torch.utils.data import DataLoader
import torch.nn.functional as F
from torchvision import transforms,datasets
import torch.optim as optim
设置数据的transform转换并设置其他的超参数 将数据转换为Tensor张量,然后设置Normalization
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,),(0.3081,))
])
batch_size = 256
num_workers = 2
device = ("cuda:0" if torch.cuda.is_available() else "cpu")
epochs = 10
获取数据集 如果没有数据集需要把download设置为True,这时就会自动从网上下载需要的数据集
tran_data = datasets.MNIST(root="./MNIST",
train=True,
transform=transform,
download=False)
test_data = datasets.MNIST(root="./MNIST",
train=False,
transform=transform,
download=False)
加载数据集 有了数据集需要将数据集里面的数据取出来,这时用的是DataLoader类
train_iter = DataLoader(tran_data,batch_size=batch_size,shuffle=True,num_workers=num_workers)
test_iter = DataLoader(test_data,batch_size=batch_size,shuffle=False,num_workers=num_workers)
定义网络结构并设置优化器
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1,6,kernel_size=5,padding=2)
self.conv2 = nn.Conv2d(6,16,kernel_size=5)
self.lin1 = nn.Linear(16*5*5,120)
self.lin2 = nn.Linear(120,84)
self.lin3 = nn.Linear(84,10)
self.avg = nn.AvgPool2d(kernel_size=2,stride=2)
def forward(self,X):
X = self.avg(F.relu(self.conv1(X)))
X = self.avg(F.relu(self.conv2(X)))
X = X.view(-1,400)
X = self.lin3(self.lin2(self.lin1(X)))
out = F.softmax(X,dim=1)
return out
model = Net().to(device)
# 使用Adam优化器
optimizer = optim.Adam(model.parameters())
进行模型的训练
def train_model(epoch):
model.train()
for batch_idx,(X,y) in enumerate(train_iter):
X,y = X.to(device),y.to(device)
y_hat = model(X)
loss = F.cross_entropy(y_hat,y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 3000 == 0:
print("Train Epoch : {} \t Loss : {:.6f}".format(epoch, loss.item()))
def test_model():
model.eval()
correct = 0.0
test_loss = 0.0
with torch.no_grad():
for X,y in test_iter:
X,y = X.to(device),y.to(device)
y_hat = model(X)
test_loss += F.cross_entropy(y_hat,y).item()
pred = y_hat.argmax(dim=1)
correct += pred.eq(y.view_as(pred)).sum().item()
test_loss /= len(test_iter.dataset)
print("Test —— Average loss : {:.4f}, Accuracy : {:.3f}\n".format(test_loss,
100.0 * correct / len(test_iter.dataset)))
def main():
for epoch in range(epochs):
train_model(epoch)
test_model()
if __name__ == '__main__':
main()
在实现的过程中发现了几个易错的bug
首先如果需要使用CPU的多线程,也就是设置num_workers>1,那么这个时候就需要使用main函数,然后再调用main函数,否则会报错 在定义网络类的时候如果没有使用super()或者init这个函数拼写错误那么就会报Adam优化器并没有接受到参数的错误。