基于 MindSpore 的图像分类模型训练
·
以下基于 MindSpore 框架,结合你提供的模型训练策略调整参数和训练步骤,给出一个完整的示例代码,并对应解释各部分与策略的关联:
示例任务:基于 MindSpore 的图像分类模型训练(以 ResNet18 为例)
python
运行
import mindspore as ms
from mindspore import nn, dataset, transforms
from mindspore.dataset import ImageFolderDataset
from mindspore.train import Model, Callback
from mindspore.train.callback import LossMonitor, ModelCheckpoint, CheckpointConfig
import mindspore.dataset.vision as vision
import mindspore.dataset.transforms as transforms
# ----------------------------
# 1. 数据读取与预处理(对应训练步骤1)
# ----------------------------
def create_dataset(data_dir, batch_size=32, is_train=True):
# 加载数据集
dataset = ImageFolderDataset(data_dir, shuffle=is_train)
# 数据预处理(含数据增强,对应策略4.c)
trans = []
if is_train:
# 训练时的数据增强:随机裁剪、翻转、归一化
trans += [
vision.RandomCrop((224, 224), padding=4), # 随机裁剪
vision.RandomHorizontalFlip(prob=0.5), # 随机水平翻转
]
# 共通处理:Resize、归一化、转换为Tensor
trans += [
vision.Resize((224, 224)),
vision.Normalize(mean=[0.485*255, 0.456*255, 0.406*255],
std=[0.229*255, 0.224*255, 0.225*255]),
vision.HWC2CHW(), # 转换通道顺序为[C, H, W]
transforms.TypeCast(ms.float32)
]
# 应用转换
dataset = dataset.map(operations=trans, input_columns="image")
dataset = dataset.map(operations=transforms.TypeCast(ms.int32), input_columns="label")
# 批次处理(对应策略4.b:Batch_size)
dataset = dataset.batch(batch_size, drop_remainder=True)
return dataset
# 加载训练集和验证集
train_dataset = create_dataset("./train_data", batch_size=32) # Batch_size=32
val_dataset = create_dataset("./val_data", batch_size=32, is_train=False)
# ----------------------------
# 2. 定义网络模型(对应训练步骤2和策略1:网络结构)
# ----------------------------
# 选择预训练模型ResNet18(与任务相关,对应策略1),并调整输出层
class CustomResNet18(nn.Cell):
def __init__(self, num_classes=10):
super(CustomResNet18, self).__init__()
# 加载预训练的ResNet18(网络结构选择)
self.backbone = nn.ResNet18(pretrained=True)
# 替换输出层(根据任务调整参数数量,对应策略1)
self.backbone.fc = nn.Dense(self.backbone.fc.in_channels, num_classes)
def construct(self, x):
return self.backbone(x)
# 初始化模型
net = CustomResNet18(num_classes=10) # 假设10分类任务
# ----------------------------
# 3. 定义损失函数(含正则化,对应策略2)
# ----------------------------
# 损失函数+L2正则化(对应策略2.a:L2正则化)
loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
# 对权重施加L2正则化(权重衰减系数weight_decay即L2正则强度)
weight_decay = 1e-4 # L2正则参数
# ----------------------------
# 4. 定义优化器(对应策略3:优化器与学习率)
# ----------------------------
# 策略3.b:学习率(采用随时间减小的函数,对应策略3.b.ii)
# 初始学习率0.01,总训练轮数100,每30轮衰减为原来的0.1
lr_scheduler = nn.cosine_decay_lr(
min_lr=0.0001, # 最小学习率
max_lr=0.01, # 初始学习率
total_step=train_dataset.get_dataset_size() * 100, # 总步数=每轮步数×总轮数
step_per_epoch=train_dataset.get_dataset_size(), # 每轮步数
decay_epoch=100 # 衰减总轮数
)
# 优化器:Adam(基于梯度下降,对应策略3.a),加入权重衰减(L2正则)
optimizer = nn.Adam(
params=net.trainable_params(),
learning_rate=lr_scheduler, # 动态学习率
weight_decay=weight_decay # L2正则化参数
)
# ----------------------------
# 5. 模型训练(对应训练步骤3)
# ----------------------------
# 自定义回调函数(对应训练步骤3.d:模型参数保存、loss打印)
class MyCallback(Callback):
def step_end(self, run_context):
# 打印每步的loss(对应步骤3.d.ii)
cb_params = run_context.original_args()
print(f"Step: {cb_params.cur_step_num}, Loss: {cb_params.net_outputs.asnumpy()}")
# 配置模型保存(对应步骤3.d.i)
config_ckpt = CheckpointConfig(save_checkpoint_steps=100, keep_checkpoint_max=5)
ckpt_callback = ModelCheckpoint(prefix="resnet18", config=config_ckpt)
# 初始化模型
model = Model(net, loss_fn=loss_fn, optimizer=optimizer, metrics={"accuracy"})
# 训练模型(对应策略4.a:epoch=100)
print("Start training...")
model.train(
epoch=100, # 总轮数
train_dataset=train_dataset,
callbacks=[LossMonitor(per_print_times=10), ckpt_callback, MyCallback()], # 回调函数
dataset_sink_mode=False
)
# ----------------------------
# 6. 策略调整关注的指标(对应策略关注指标)
# ----------------------------
# 验证模型(获取评价指标accuracy、loss)
metrics = model.eval(val_dataset)
print(f"Validation metrics: {metrics}")
# 计算推理时间(对应指标3:推理时间)
import time
start_time = time.time()
model.predict(val_dataset.take(1)) # 推理1个batch
infer_time = time.time() - start_time
print(f"Inference time for 1 batch: {infer_time:.4f}s")
代码与策略的对应解释
1. 网络结构调整(策略 1)
- 选择预训练的 ResNet18 作为基础模型(与图像分类任务匹配),并替换输出层以适应自定义分类数(调整参数数量)。
- 若模型过拟合 / 欠拟合,可进一步调整:如增减层数(改用 ResNet34/ResNet10)、增减模块(添加注意力模块)。
2. 正则化(策略 2)
- 损失函数中通过优化器的
weight_decay参数实现 L2 正则化(weight_decay=1e-4),抑制过拟合。 - 若需添加 Dropout,可在网络中插入
nn.Dropout(keep_prob=0.5)层(对应策略 2.b)。
3. 优化器与学习率(策略 3)
- 优化器选择 Adam(基于梯度下降的自适应优化器,平衡模型稳定与准确率,对应策略 3.a)。
- 学习率采用余弦衰减策略(随训练步数增加逐渐减小,对应策略 3.b.ii),相比固定学习率更易收敛。
4. 数据相关参数(策略 4)
epoch=100:训练总轮数(策略 4.a),控制模型对数据的学习次数。batch_size=32:每次迭代的样本数(策略 4.b),影响训练效率和模型稳定性。- 数据增强:训练时添加随机裁剪、翻转(策略 4.c),扩充数据多样性以提升泛化能力。
5. 训练步骤(对应训练步骤)
- 数据读取与预处理:通过
create_dataset函数完成,转换为 MindSpore 的 Dataset 格式。 - 定义网络模型:
CustomResNet18类封装了调整后的 ResNet18 结构。 - 模型训练:
model.train接口整合了模型(net)、损失函数(loss)、优化器(opt),并通过回调函数实现参数保存和 loss 打印。
6. 关注指标(对应策略关注指标)
- 训练中通过
LossMonitor和自定义回调打印loss,验证时计算accuracy(指标 1)。 - 模型参数数量可通过
net.parameters_and_names()统计(指标 2)。 - 推理时间通过计时函数计算(指标 3),训练时间可通过训练总耗时统计(指标 4)。
通过调整上述代码中的参数(如batch_size、epoch、学习率策略、正则化系数等),可实现模型训练策略的优化,提升模型性能。
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。
更多推荐


所有评论(0)