import os
import numpy as np
import pandas as pd
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
---------------------- 1. 配置参数 ----------------------
数据集路径（需替换为你的CASIA-HWDB图像文件夹路径）
DATA_ROOT = "path/to/your/casia_hwdb_images"
汉字标签文件（需包含image_name与label的映射，格式：文件名,汉字）
LABEL_FILE = "path/to/your/labels.csv"
训练参数
BATCH_SIZE = 64
EPOCHS = 20
LEARNING_RATE = 1e-3
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
NUM_CLASSES = 3755 # CASIA-HWDB常用汉字数量
---------------------- 2. 自定义数据集类 ----------------------
class ChineseHandwritingDataset(Dataset):
def init(self, data_root, label_file, transform=None):
self.data_root = data_root
self.labels = pd.read_csv(label_file)
# 汉字到索引的映射（将汉字转为模型可识别的数字）
self.char_to_idx = {char: i for i, char in enumerate(self.labels["label"].unique())}
self.transform = transform
def __len__(self):return len(self.labels)def __getitem__(self, idx):# 读取图像img_name = self.labels.iloc[idx]["image_name"]img_path = os.path.join(self.data_root, img_name)image = Image.open(img_path).convert("L") # 转为灰度图（单通道）# 读取标签并转为数字label_char = self.labels.iloc[idx]["label"]label = self.char_to_idx[label_char]# 图像预处理if self.transform:image = self.transform(image)return image, label
---------------------- 3. 定义CNN模型（轻量级） ----------------------
class HandwritingCNN(nn.Module):
def init(self, num_classes=NUM_CLASSES):
super(HandwritingCNN, self).init()
# 卷积层：提取图像特征
self.conv_layers = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, padding=1), # 输入1通道，输出32通道
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2), # 下采样：尺寸减半
nn.Conv2d(32, 64, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),nn.Conv2d(64, 128, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2))# 全连接层：分类self.fc_layers = nn.Sequential(nn.Flatten(), # 展平卷积输出（需根据输入图像尺寸调整后续维度）nn.Linear(128 * 16 * 16, 1024), # 假设输入图像为128x128，经3次下采样后为16x16nn.ReLU(),nn.Dropout(0.5), # 防止过拟合nn.Linear(1024, num_classes))def forward(self, x):x = self.conv_layers(x)x = self.fc_layers(x)return x
---------------------- 4. 数据加载与预处理 ----------------------
图像预处理：统一尺寸、转为Tensor、标准化
transform = transforms.Compose([
transforms.Resize((128, 128)), # 统一图像尺寸为128x128
transforms.ToTensor(), # 转为Tensor（维度：C×H×W）
transforms.Normalize(mean=[0.5], std=[0.5]) # 灰度图标准化（单通道均值/标准差）
])
加载数据集（此处简化为单数据集，实际需拆分train/val）
dataset = ChineseHandwritingDataset(
data_root=DATA_ROOT,
label_file=LABEL_FILE,
transform=transform
)
拆分训练集与验证集（8:2）
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
数据加载器（批量读取+打乱）
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
---------------------- 5. 模型训练与验证 ----------------------
def train_model():
# 初始化模型、损失函数、优化器
model = HandwritingCNN(num_classes=NUM_CLASSES).to(DEVICE)
criterion = nn.CrossEntropyLoss() # 分类任务常用损失
optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4) # 带权重衰减的Adam
# 训练循环
for epoch in range(EPOCHS):model.train() # 训练模式（启用Dropout）train_loss = 0.0train_correct = 0# 训练批次for images, labels in train_loader:images, labels = images.to(DEVICE), labels.to(DEVICE)# 前向传播outputs = model(images)loss = criterion(outputs, labels)# 反向传播+优化optimizer.zero_grad() # 清空梯度loss.backward() # 计算梯度optimizer.step() # 更新参数# 统计训练损失与准确率train_loss += loss.item() * images.size(0)_, preds = torch.max(outputs, 1) # 获取预测结果（概率最大的类别）train_correct += torch.sum(preds == labels.data)# 计算训练集平均损失与准确率train_loss = train_loss / len(train_loader.dataset)train_acc = train_correct.double() / len(train_loader.dataset)# 验证模型model.eval() # 验证模式（关闭Dropout）val_loss = 0.0val_correct = 0with torch.no_grad(): # 禁用梯度计算（节省内存）for images, labels in val_loader:images, labels = images.to(DEVICE), labels.to(DEVICE)outputs = model(images)loss = criterion(outputs, labels)val_loss += loss.item() * images.size(0)_, preds = torch.max(outputs, 1)val_correct += torch.sum(preds == labels.data)val_loss = val_loss / len(val_loader.dataset)val_acc = val_correct.double() / len(val_loader.dataset)# 打印每轮结果print(f"Epoch [{epoch+1}/{EPOCHS}]")print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}\n")# 保存训练好的模型
torch.save(model.state_dict(), "chinese_handwriting_model.pth")
print("模型已保存为：chinese_handwriting_model.pth")
---------------------- 6. 启动训练 ----------------------
if name == "main":
train_model()
