Google Coral上的图像分割模型训练与推理工作流求助

一、训练前准备 #

• 环境配置：安装TensorFlow 2.x及相关依赖包，执行命令：

  pip install tensorflow pillow tflite-support pycoral
  

• 数据集准备：
  • 采用语义分割标准格式：每张输入图像对应一张同尺寸的掩码图，掩码图中每个像素值代表对应类别索引
  • 推荐将数据集转换为TFRecord格式提升训练效率，或直接用ImageDataGenerator加载图像与掩码对

二、模型训练 #

1. 选择适配Coral的模型架构 #

优先选用轻量级、量化友好的架构，比如基于MobileNetV2的DeepLabV3，或自定义轻量U-Net，避免复杂大模型（如带VGG16 backbone的原始U-Net）。

2. 训练代码示例 #

import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Input, Conv2DTranspose, concatenate, Conv2D
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 构建轻量级分割模型
def build_lightweight_segmenter(input_shape=(256,256,3), num_classes=2):
    base_model = MobileNetV2(input_shape=input_shape, include_top=False, weights='imagenet')
    # 提取多尺度特征层
    feature_layers = [
        'block_1_expand_relu',  # 64x64特征图
        'block_3_expand_relu',  # 32x32特征图
        'block_6_expand_relu',  # 16x16特征图
        'block_13_expand_relu', # 8x8特征图
        'block_16_project'      # 4x4特征图
    ]
    features = [base_model.get_layer(name).output for name in feature_layers]

    # 上采样与特征融合
    x = features[-1]
    for feat in reversed(features[:-1]):
        x = Conv2DTranspose(256, (2,2), strides=(2,2), padding='same')(x)
        x = concatenate([x, feat])
    # 输出类别预测
    x = Conv2D(num_classes, (1,1), activation='softmax')(x)

    model = Model(inputs=base_model.input, outputs=x)
    return model

# 配置数据集生成器（示例框架，需自行适配实际数据集结构）
train_datagen = ImageDataGenerator(rescale=1./255)
train_generator = train_datagen.flow_from_directory(
    'train_data_dir',
    target_size=(256,256),
    class_mode='input'
)

# 编译并训练模型
model = build_lightweight_segmenter(num_classes=2)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_generator, epochs=20, validation_split=0.2)
model.save('custom_segmenter.h5')

3. 量化准备 #

训练时尽量使用Coral支持的标准算子，避免自定义层；若追求更高精度，可开启量化感知训练（在模型构建时插入量化节点），后续转换时精度损失更小。

三、模型转换为Coral兼容格式 #

1. 转换为INT8量化的TensorFlow Lite模型 #

converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 启用后训练量化
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 提供校准数据集用于量化校准
def representative_dataset():
    for batch in train_generator.take(100):
        yield [batch[0]]
converter.representative_dataset = representative_dataset
# 指定INT8输入输出
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

# 生成量化模型
tflite_quant_model = converter.convert()
with open('quant_segmenter.tflite', 'wb') as f:
    f.write(tflite_quant_model)

2. 转换为Edge TPU专用模型 #

使用Coral提供的edgetpu_compiler工具转换：

edgetpu_compiler quant_segmenter.tflite

执行完成后会生成quant_segmenter_edgetpu.tflite，这就是可在Coral设备上运行的模型。

四、Coral设备上的推理 #

from pycoral.utils import edgetpu
from pycoral.adapters import common, segment
from PIL import Image
import numpy as np

# 加载Edge TPU模型
interpreter = edgetpu.make_interpreter('quant_segmenter_edgetpu.tflite')
interpreter.allocate_tensors()

# 预处理输入图像
input_size = common.input_size(interpreter)
image = Image.open('test_img.jpg').resize(input_size)
input_data = np.array(image, dtype=np.int8)
common.set_input(interpreter, input_data)

# 运行推理
interpreter.invoke()

# 获取分割结果
segment_mask = segment.get_output(interpreter)
# 转换为类别掩码（每个像素对应类别索引）
class_mask = np.argmax(segment_mask, axis=-1)

# 可视化结果
mask_img = Image.fromarray(class_mask.astype(np.uint8) * 100)  # 用灰度值区分类别
mask_img.show()

关键注意事项 #

• 算子兼容性：避免使用Edge TPU不支持的算子（如tf.nn.softmax_v2、部分自定义损失函数），可参考Coral官方算子支持列表排查
• 精度优化：若后训练量化精度不足，改用量化感知训练，即在训练过程中模拟量化操作
• 数据集规范：掩码图必须与输入图像尺寸严格一致，类别索引需连续无间隔

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