如何调用U-Net训练60轮后的权重实现CT图像预测？

一、修改训练代码，保存中间轮次权重（推荐方案） #

原训练代码未配置中间权重保存逻辑，需添加ModelCheckpoint回调实现轮次权重保存，步骤如下：

1. 导入所需模块

from tensorflow.keras.callbacks import ModelCheckpoint

2. 更新回调函数列表，添加权重保存规则

callbacks = [
    tf.keras.callbacks.TensorBoard(
        log_dir='logs',
        histogram_freq=1,
        write_graph=True,
        write_images=True,
    ),
    ModelCheckpoint(
        filepath='model_weights_epoch_{epoch:03d}.h5',  # 文件名包含轮次编号
        save_weights_only=True,  # 仅保存权重，减少存储占用
        save_freq='epoch'  # 每轮训练结束后保存一次
    )
]

3. 重新运行训练代码，完成后会生成如model_weights_epoch_060.h5的第60轮权重文件。

4. 加载第60轮权重进行预测

# 构建与训练时一致的模型结构
model = UNet_model_2D(256, 256, 1)

# 加载第60轮的权重文件
model.load_weights('model_weights_epoch_060.h5')

# 执行预测流程
X_test = nib.load('test_MRIs.nii.gz').get_fdata()
predicted_data = model.predict(X_test, verbose=1)
image = nib.Nifti1Image(predicted_data, affine=None)
nib.save(image, 'predicted_sCTs_epoch60.nii.gz')

二、无中间权重时的补救方案 #

如果已经完成200轮训练且未保存中间权重，只能重新训练：

• 修改训练代码的epochs参数为60，重新训练至第60轮后，直接保存模型：

# 训练时设置epochs=60
results = model.fit(
    x=X,
    y=Y,
    batch_size=16, 
    epochs=60,  # 只训练到第60轮
    verbose=1,
    callbacks=callbacks,
    validation_data = (x_val, y_val),
)

# 保存第60轮的模型
model.save('pCT_2D_deep_large_batch16_epoch60', save_format='tf')

之后直接加载该模型进行预测即可。

内容的提问来源于stack exchange，提问作者user20874428