TensorPlay AI

CNN 图像分类

两层卷积、两层全连接、DataLoader、交叉熵与 Adam，组成完整分类训练路径。

模型与训练

import tensorplay as tp
import tensorplay.nn as nn
import tensorplay.optim as optim
from tensorplay.utils.data import DataLoader
from tensorplay.datasets import MNIST

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.flatten(1)
        return self.fc2(self.relu(self.fc1(x)))

model = CNN()
train_dataset = MNIST(root='./data', train=True, download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(5):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

评估

切换到 eval 模式，在 tp.no_grad() 中遍历测试集，取最大类别并计算正确率。

Ask DeepWiki