Mask R-CNN训练Epoch 1/5卡住无进度问题求助

训练卡在Epoch 1/5无进度，核心是数据集加载代码存在多处致命逻辑错误，导致数据生成器无法返回有效训练样本，具体问题和修复方案如下：

已定位的核心问题 #

• 边界框提取逻辑完全失效：extract_boxes函数遍历XML标注的bndbox节点时，计算完坐标coors后没有将其追加到boxes列表，导致返回的boxes永远是空列表，后续生成mask时维度为0，训练流程直接卡在数据读取环节。
• 数据集划分逻辑不匹配现有数据量：原教程的划分规则是针对千张级数据集设计（训练集取id<150、验证集取id>=150），你当前只有32张图片，按这个规则验证集样本数为0，训练启动时加载验证集会直接卡住。
• 类名读取存在格式问题：读取class_list.txt时用readlines()会保留每行末尾的换行符，导致注册的类名带\n后缀，后续匹配Car/Bicycle类名时会触发索引错误。
• 类别判断逻辑硬编码：load_mask里写死了第一个框是Car、其余是Bicycle，还把读取XML中类别name字段的代码注释了，就算能跑也会出现类别完全对应错误的问题。
• 训练步长配置不合理：STEPS_PER_EPOCH设为100、VALIDATION_STEPS设为50，远大于你现有数据集的样本量，会导致数据生成器无意义空转。

修复后的代码 #

1. 数据集类修正代码 #

class ModelDataset(Dataset):
    # 加载数据集定义
    def load_dataset(self, dataset_dir, is_train=True):
        # 读取类别，去掉每行末尾换行符
        lines = []
        with open('OpenLabeling-master/main/class_list.txt') as f:
            lines = [line.strip() for line in f.readlines()]
        count = 1
        for line in lines:
            self.add_class("dataset", count, line)
            count += 1
        # 定义数据路径
        images_dir = dataset_dir + '/images/'
        annotations_dir = dataset_dir + '/annots/'
        # 收集所有有效图片id，按8:2拆分训练/验证集适配小数据集
        all_image_ids = []
        for filename in listdir(images_dir):
            image_id = filename[:-4]
            if image_id in ['00090']:
                continue
            # 跳过非数字命名文件防止类型报错
            if not image_id.isdigit():
                continue
            all_image_ids.append(image_id)
        # 拆分数据集
        split_idx = int(len(all_image_ids)*0.8)
        train_ids = all_image_ids[:split_idx]
        val_ids = all_image_ids[split_idx:]
        use_ids = train_ids if is_train else val_ids
        
        for image_id in use_ids:
            img_path = images_dir + image_id + '.jpg'
            ann_path = annotations_dir + image_id + '.xml'
            self.add_image('dataset', image_id=image_id, path=img_path, annotation=ann_path)
 
    # 加载单张图的所有边界框
    def extract_boxes(self, filename):
        root = ElementTree.parse(filename)
        boxes = list()
        box_class_list = list()
        # 同时提取每个边界框坐标和对应类别
        for obj in root.findall('.//object'):
            cls_name = obj.find('name').text.strip()
            box = obj.find('bndbox')
            xmin = int(box.find('xmin').text)
            ymin = int(box.find('ymin').text)
            xmax = int(box.find('xmax').text)
            ymax = int(box.find('ymax').text)
            coors = [xmin, ymin, xmax, ymax]
            boxes.append(coors) # 补全之前缺失的坐标追加操作
            box_class_list.append(cls_name)
        # 提取图片尺寸
        width = int(root.find('.//size/width').text)
        height = int(root.find('.//size/height').text)
        return boxes, box_class_list, width, height
 
    # 加载单张图的mask
    def load_mask(self, image_id):
        info = self.image_info[image_id]
        path = info['annotation']
        boxes, box_cls_names, w, h = self.extract_boxes(path)
        # 生成mask数组
        masks = zeros([h, w, len(boxes)], dtype='uint8')
        class_ids = list()
        for i in range(len(boxes)):
            box = boxes[i]
            row_s, row_e = box[1], box[3]
            col_s, col_e = box[0], box[2]
            masks[row_s:row_e, col_s:col_e, i] = 1
            # 动态匹配类别，取消硬编码
            class_ids.append(self.class_names.index(box_cls_names[i]))
        return masks, asarray(class_ids, dtype='int32')

2. 配置参数修正 #

调整步长参数适配32张小数据集，其余参数保持不变即可（显存不足时可将backbone换成resnet50）：

# 训练集约25张样本，步长设为对应值
STEPS_PER_EPOCH = 25
# 验证集约7张样本，步长设为对应值
VALIDATION_STEPS = 5

3. 训练启动代码（无需修改，确认数据集路径正确即可） #

model = MaskRCNN(mode='training', model_dir='./', config=config)
model.load_weights('mask_rcnn_coco.h5', by_name=True, exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",  "mrcnn_bbox", "mrcnn_mask"])
model.train(train_set, test_set, learning_rate=config.LEARNING_RATE, epochs=5, layers='heads')

启动前验证 #

改完代码后先单独实例化数据集，调用load_dataset后执行prepare()，打印train_set.num_images确认训练集、验证集图片数量正常，再随机取1张图调用load_mask确认返回的mask和class_ids维度不为空，再启动训练就不会卡住。

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