知识点回顾：
- 彩色和灰度图片测试和训练的规范写法：封装在函数中
- 展平操作：除第一个维度batchsize外全部展平
- dropout操作：训练阶段随机丢弃神经元，测试阶段eval模式关闭dropout
作业：仔细学习下测试和训练代码的逻辑，这是基础，这个代码框架后续会一直沿用，后续的重点慢慢就是转向模型定义阶段了。
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
def load_data(color_mode='rgb', batch_size=32):
"""
加载彩色或灰度图片数据集，并进行规范化处理
Args:
color_mode: 'rgb'或'grayscale'，指定图片颜色模式
batch_size: 批次大小
Returns:
train_loader: 训练数据加载器
test_loader: 测试数据加载器
"""
# 定义数据转换
if color_mode == 'rgb':
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
else: # 灰度图
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
])
# 加载数据集
train_dataset = datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
return train_loader, test_loader
def flatten_tensor(x):
"""
对输入张量进行展平操作，保留第一个维度(batch_size)
Args:
x: 输入张量，形状为(batch_size, ...)
Returns:
展平后的张量，形状为(batch_size, num_features)
"""
return x.view(x.size(0), -1)
class SimpleModel(nn.Module):
"""简单模型，包含展平和dropout操作"""
def __init__(self, input_size, hidden_size, num_classes, dropout_rate=0.5):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.dropout = nn.Dropout(dropout_rate)
self.fc2 = nn.Linear(hidden_size, num_classes)
self.relu = nn.ReLU()
def forward(self, x):
# 展平操作
x = flatten_tensor(x)
x = self.fc1(x)
x = self.relu(x)
# Dropout操作（训练时随机丢弃，测试时关闭）
x = self.dropout(x)
x = self.fc2(x)
return x
def train_model(model, train_loader, criterion, optimizer, device, epochs=5):
"""
训练模型，包含dropout操作
Args:
model: 待训练的模型
train_loader: 训练数据加载器
criterion: 损失函数
optimizer: 优化器
device: 计算设备
epochs: 训练轮数
"""
model.train() # 启用训练模式（dropout生效）
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader):.4f}')
def test_model(model, test_loader, device):
"""
测试模型，关闭dropout
Args:
model: 待测试的模型
test_loader: 测试数据加载器
device: 计算设备
Returns:
accuracy: 模型准确率
"""
model.eval() # 启用评估模式（dropout关闭）
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f'Accuracy: {accuracy:.2f}%')
return accuracy
# 示例用法
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载彩色图片数据
train_loader, test_loader = load_data(color_mode='rgb', batch_size=32)
# 计算输入大小
sample_data, _ = next(iter(train_loader))
input_size = sample_data.view(sample_data.size(0), -1).size(1)
# 创建模型
model = SimpleModel(input_size, 128, 10, dropout_rate=0.5).to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练模型
train_model(model, train_loader, criterion, optimizer, device, epochs=5)
# 测试模型
test_model(model, test_loader, device) Epoch 1/5, Loss: 1.5823
Epoch 2/5, Loss: 1.2356
Epoch 3/5, Loss: 1.0432
Epoch 4/5, Loss: 0.9124
Epoch 5/5, Loss: 0.8231
Accuracy: 68.32%