第五届计图jittor人工智能挑战赛-超声图像的智能筛查与分级（赛道一）

简介

本项目包含了第五届计图jittor人工智能挑战赛-超声图像的智能筛查与分级（赛道一）的代码实现。

首先，使用提供的baseline在训练集上进行训练，发现在验证集上可以收敛到0.95左右的准确率，将训练好的模型在A榜的测试集上进行测试发现，准确率为0.78左右，说明训练集和测试集的数据部分差距较大。

其次，可视化了训练集和验证集的类别标签分布如下图所示，发现训练集和测试集的标签分布基本一致，但是类别不平衡，类别3、4非常少。

因此，将目标定位为缓和类别不平衡问题和提高模型的泛化性能。主要采取了以下三个策略：

• TrivialAugment进行数据增强：通过对训练数据进行各种随机变换（如旋转、裁剪、翻转、颜色抖动等），生成“新”的训练样本，从而扩充训练集的数据多样性，让模型在训练时看到更多样化的输入，迫使它学习更具鲁棒性的特征，而不是记住训练集的特定模式（过拟合）。
• SAM优化器：不仅最小化训练损失，还同时最小化损失函数的“尖锐度”（sharpness），即在参数空间中寻找平坦的极小值点（flat minima），经验表明，平坦的损失极小值（周围参数小扰动对损失影响小）通常对应更好的泛化性能。
• Label Smoothing：将硬标签（hard labels，如 [0, 0, 1]）替换为软标签（soft labels，如 [0.1, 0.1, 0.8]），避免模型对某个类别过于自信。

安装

本项目可在 1 张 4090 上运行，训练时间约为 0.1 小时。

运行环境

• ubuntu 22.04 LTS
• python >= 3.7
• jittor >= 1.3.0

安装依赖

使用conda从jittor_env.yml创建环境

conda env create -f jittor_env.yml

超参数管理

使用hydra管理超参数，可以在config.yaml中进行更改，超参数说明如下：

test_model：是否测试模型，测试设为True，训练设为False
train_data_dir：训练数据的目录
train_label_path：训练标签的路径
val_data_dir：验证数据的目录
val_label_path：验证标签的路径
model_dir：保存训练好的模型的路径
test_data_dir：测试数据的目录
model_ckpt：测试时要加载的模型权重的路径
result_path：测试结果的保存路径

训练

在config.yaml文件设置test_model=False，并将对应参数指向相应路径后，运行：

python main.py

推理

在config.yaml文件设置test_model=True，并将对应参数指向相应路径后，运行：

python main.py

致谢

本项目参考了NIPS2023的文章：Simplifying Neural Network Training Under Class Imbalance。