【PyTorch】训练一个最简单的CNN
导入相关包
torch.nn.functional
中包含relu()
,maxpool2d()
等 CNN 常用操作。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
显示 pytorch 环境版本及是否使用 GPU
print(torch.__version__)
print(torchvision.__version__)
print(torch.cuda.is_available())
下载FashionMNIST数据集并转换格式为Tensor(Extract & Transform)
train_set = torchvision.datasets.FashionMNIST(
root='./data/FashionMNIST'
,train=True
,download=True
,transform=transforms.Compose([
transforms.ToTensor()
])
)
继承nn.Moudle
建立模型,主要是写forward
class Network(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
self.fc1 = nn.Linear(in_features=12*4*4, out_features=120)
self.fc2 = nn.Linear(in_features=120, out_features=60)
self.out = nn.Linear(in_features=60, out_features=10)
def forward(self,t):
#(1) input layer
t = t
#(2) hidden conv layer
t = self.conv1(t)
t = F.relu(t)
t = F.max_pool2d(t, kernel_size=2, stride=2)
#(3) hidden conv layter
t = self.conv2(t)
t = F.relu(t)
t = F.max_pool2d(t, kernel_size=2, stride=2)
#(4) hidden linear layer
t = t.reshape(-1, 12*4*4)
t = self.fc1(t)
t = F.relu(t)
#(5) hidden linear layer
t = self.fc2(t)
t = F.relu(t)
#(6) output layer
t = self.out(t)
return t
写工具函数:获得本次预测结果的正确个数
def get_num_correct(preds, labels):
return preds.argmax(dim=1).eq(labels).sum().item()
定义网络,优化器,加载(Load)数据(设定 batch 大小等)
network = Network()
optimizer = optim.Adam(network.parameters(), lr=0.01)
train_loader = torch.utils.data.DataLoader(
train_set
,batch_size=100
,shuffle=True
)
训练
for epoch in range(10):
total_loss = 0
total_correct = 0
for batch in train_loader: # Get Batch
images, labels = batch
preds = network(images) # Pass Batch
loss = F.cross_entropy(preds, labels) # Calculate Loss
optimizer.zero_grad()
loss.backward() # Calculate Gradients
optimizer.step() # Update Weights
total_loss += loss.item()
total_correct += get_num_correct(preds, labels)
print(
"epoch", epoch,
"total_correct:", total_correct,
"loss:", total_loss
)