记一次模型训练——基于神经网络的PBR纹理分类

任务描述

给定图片识别出是BaseColor/Metallic/…等类型。

模型

• ResNet 34
• ResNet 50

输入大小：512×512×3 训练精度：FP16

数据集

MatSynth

DatasetDict(
	train: Dataset({
		features: ['name', 'category', 'metadata', 'basecolor', 'diffuse', 'displacement', 'height', 'metallic', 'normal', 'opacity', 'roughness', 'specular', 'blend_mask'],
		num_rows: 5700
	})
	test: Dataset({
		features: ['name', 'category', 'metadata', 'basecolor', 'diffuse', 'displacement', 'height', 'metallic', 'normal', 'opacity', 'roughness', 'specular', 'blend_mask'],
		num_rows: 89
	})
)

选取的数据

• basecolor
• metallic
• roughness
• normal
• displacement (displacement + height)

数据预处理

数据清理

1. 纹理通道数问题

• 统一用三通道
• 对于四通道，直接舍弃 Alpha 通道；
• 对于单通道，拓展通道至三通道，复制原通道数据

2. 纯色纹理问题 纯色纹理直接记为无效数据，训练前先扫描有效数据数量。

3. 类别不均衡问题 根据各类别的通道数和有效数设置损失函数加类别权重。

数据采样

1. 随机Crop 从图片上随机 Crop 从256×256至1024×1024随机大小的分块，再 Resize 到 512×512 输入模型。

2. 随机旋转 对 Crop 窗口随机旋转，但不能裁剪到超出图像大小的内容。

3. Shuffle 由于 MatSynth 数据集是一个样本 (Row) 对应一套纹理 (Col)，而当前任务是识别纹理类型，故不需要关注一套纹理 (1 Row) 内部的关系，加上类别不均衡，所以这里的 Shuffle 不能仅仅打乱数据集样本的顺序 (Row 顺序)，而是直接打乱数据集中所有纹理，甚至包括被拆分的三通道纹理，每个通道作为独立的纹理来打乱。

训练性能优化

CPU: AMD Ryzen 3 PRO 4350G RAM: 64GB GPU: Tesla P100 16GB × 1 存储: 机械硬盘 RAID0 (1TB × 2) 由于当前机器的存储性能非常差，加载数据集时会有非常高IO延迟，故需要做以下优化：

1. 同时串行训练多个模型 由于IO延迟较高，分别训练多个模型时需要多次等待IO，故将多个模型的训练放在一起，即读一次IO，串行训练多个模型，训练时保持三个模型在内存中。

2. 内存盘与分区 Shuffle 当前机器的存储无法支撑高强度随机读写，将使用内存盘存储部分数据集以加快数据加载时间，具体流程为：每次随机选定多个 Shard (Arrow)，复制到内存盘中并创建新的dataset_info.json，加载内存盘中的数据集，执行 Shuffle，进入训练。当 Shard 耗尽即完成一个 Epoch。

3. 手动实现 Arrow 解码 HuggingFace Arrow 使用的是 PIL，实测下来解码速度很慢，应当绕过 PIL，直接从 Arrow 表读取原始字节，再用 cv2 解码。

4. 异步加载和训练 开启两个线程，一个专门复制文件、加载数据，一个专门执行训练模型，具体流程为：线程1加载数据 → 线程2发现可用数据消费数据进行训练 → 数据被消费线程1加载数据 → …

5. 断点续训

训练参数与指标

Batch Size: 64 Learn Rate: 2.5e-5 Scheduler: CosineAnnealingLR (Per Window) Epochs: 5

ResNet-34

• 平均损失:
  • 训练集 0.2152
  • 验证集 0.1264
• 平均正确率:
  • 训练集 95.15%
  • 验证集 95.20%

ResNet50

• 平均损失:
  • 训练集 0.2322
  • 验证集 0.1586
• 平均正确率:
  • 训练集 94.30%
  • 验证集 97.52%