鏖战mask2former(二)
事情变得越发麻烦起来...
感觉算是经验帖,以后可能可以作为自己DEBUG的范式了。
因为和实际工程挂钩,所以还暂时不能完整公开干了什么,就尽可能以替代的方式描述。后续如果论文发了中了再补个论文的链接。
前言
前情提要:
按照开发周期(一周)来进行记录,(二)记录的是第二周的内容。
请善用导航,本文内容特别长。
BEFORE EVERYTHING GET STARTED
需求描述
也算是敏捷开发的一个迭代周期了。
这一部分主要开发时间:2025-04-13至2025-04-19。
用户故事:作为一个计算机视觉开发工程师(?我不是我没有),我希望能检出并修复训练中的问题,以让新模型的训练效果与正常水平齐平。
这是一条难度非常大的用户故事,如果是我肯定会分配更多的故事点,但是时间不允许了。
更具体些
喜报:AP=0!
这个结果意味着,除了流程确实跑通了以外,整个代码的实现中存在着较多极为严重的问题,这使得模型几乎没有学习到任何内容。
考虑到在上一周中,改动的代码包括了Mapper、Meta-Arch、Decoder、Criterion四块组件,因此着重要检查这四块组件的问题。
事实上,上一周基本上只是粗略的看了一遍代码,除了检查过一部分的Meta-Arch,其余部分的继承实现几乎全是让AI写的。
这是一个极度危险的行为,也是导致了这一周工作量巨大的罪魁祸首。这并不是好的AI应用实践。
具体的事项包括了:
- 检查原始模型的实现
- 检查代码中对于预处理部分的实现
- 检查backbone阶段输出层与decoder部分的链接实现
- 检查多损失实现
- 检查pretrained
- 检查、修改继承方法
- 测试
LET’S DIVE INTO CODE
预处理阶段
预处理阶段的内容被配置在了mask2former/data/dataset_mappers
,以COCO为例:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def build_transform_gen(cfg, is_train):
"""
Create a list of default :class:`Augmentation` from config.
Now it includes resizing and flipping.
Returns:
list[Augmentation]
"""
assert is_train, "Only support training augmentation"
image_size = cfg.INPUT.IMAGE_SIZE
min_scale = cfg.INPUT.MIN_SCALE
max_scale = cfg.INPUT.MAX_SCALE
augmentation = []
if cfg.INPUT.RANDOM_FLIP != "none":
augmentation.append(
T.RandomFlip(
horizontal=cfg.INPUT.RANDOM_FLIP == "horizontal",
vertical=cfg.INPUT.RANDOM_FLIP == "vertical",
)
)
augmentation.extend([
T.ResizeScale(
min_scale=min_scale, max_scale=max_scale, target_height=image_size, target_width=image_size
),
T.FixedSizeCrop(crop_size=(image_size, image_size)),
])
return augmentation
代码的内容可以描述为,仅在Train阶段执行数据增强,包括了RandomFlip和Resize-Crop两步骤,使得最终的尺寸为image_size=[1024, 1024]。
而在AI写的继承版本(继承自DatasetMapper),则忽视了数据的预处理,相当于直接把原始尺寸扔进模型了,存在较多隐患,在这一部分上需要做出修改。
backbone与decoder的集成
在原始的模型文件maskformer_model.py
中,有如下内容:
1
2
3
4
5
6
7
8
9
10
@META_ARCH_REGISTRY.register()
class MaskFormer(nn.Module):
# ...
def forward(self, batched_inputs):
# ...
features = self.backbone(images.tensor)
outputs = self.sem_seg_head(features)
# ...
其中,sem_seg_head追溯到mask_former_head.py
:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
@SEM_SEG_HEADS_REGISTRY.register()
class MaskFormerHead(nn.Module):
# ...
def forward(self, features, mask=None):
return self.layers(features, mask)
def layers(self, features, mask=None):
mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
if self.transformer_in_feature == "multi_scale_pixel_decoder":
predictions = self.predictor(multi_scale_features, mask_features, mask)
# ...
return predictions
再进一步追溯到multiscale pixel-decoder的实现:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
@SEM_SEG_HEADS_REGISTRY.register()
class MSDeformAttnPixelDecoder(nn.Module):
# ...
@configurable
def __init__(
self,
input_shape: Dict[str, ShapeSpec],
*,
transformer_dropout: float,
transformer_nheads: int,
transformer_dim_feedforward: int,
transformer_enc_layers: int,
conv_dim: int, # 注意这个参数
mask_dim: int,
norm: Optional[Union[str, Callable]] = None,
# deformable transformer encoder args
transformer_in_features: List[str],
common_stride: int,
):
# ...
lateral_convs = []
output_convs = []
use_bias = norm == ""
for idx, in_channels in enumerate(self.feature_channels[:self.num_fpn_levels]):
lateral_norm = get_norm(norm, conv_dim)
output_norm = get_norm(norm, conv_dim)
lateral_conv = Conv2d(
in_channels, conv_dim, kernel_size=1, bias=use_bias, norm=lateral_norm
)
output_conv = Conv2d(
conv_dim,
conv_dim,
kernel_size=3,
stride=1,
padding=1,
bias=use_bias,
norm=output_norm,
activation=F.relu,
)
weight_init.c2_xavier_fill(lateral_conv)
weight_init.c2_xavier_fill(output_conv)
self.add_module("adapter_{}".format(idx + 1), lateral_conv)
self.add_module("layer_{}".format(idx + 1), output_conv)
lateral_convs.append(lateral_conv)
output_convs.append(output_conv)
# Place convs into top-down order (from low to high resolution)
# to make the top-down computation in forward clearer.
self.lateral_convs = lateral_convs[::-1]
self.output_convs = output_convs[::-1]
# ...
@autocast(enabled=False)
def forward_features(self, features):
# ..
# append `out` with extra FPN levels
# Reverse feature maps into top-down order (from low to high resolution)
for idx, f in enumerate(self.in_features[:self.num_fpn_levels][::-1]):
x = features[f].float()
lateral_conv = self.lateral_convs[idx] # 注意这里
output_conv = self.output_convs[idx]
cur_fpn = lateral_conv(x)
# Following FPN implementation, we use nearest upsampling here
y = cur_fpn + F.interpolate(out[-1], size=cur_fpn.shape[-2:], mode="bilinear", align_corners=False)
y = output_conv(y)
out.append(y)
for o in out:
if num_cur_levels < self.maskformer_num_feature_levels:
multi_scale_features.append(o)
num_cur_levels += 1
return self.mask_features(out[-1]), out[0], multi_scale_features
而通过观察原始backbone的两个实现,来自detectron2框架内部的resnet实现和swin实现,发现out_features是并不匹配的。
resnet的out_features=[256, 512, 1024, 2048],而swin的out_features=[96, 192, 384, 768]。通过对上述代码的分析,其又经过了lateral_conv和output_conv的步骤,使得最终的每一层结果被固定转换到conv_dim=256,再进入后续的transformer decoder流程。
那么其对于输入,理论上只要能匹配上res2~res5的多级结构就应该可以完成后续的训练,不需要任额外的转换;而我自己的代码却做了强制转换到features=112才能运行,而且效果相对较差,需要排查其他原因的问题。
如果你比我对数字敏感的话,那么这里你应该可以看出来问题出在哪里。
往下看,看看你猜对了没:
报错点在以下位置:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
@SEM_SEG_HEADS_REGISTRY.register()
class MSDeformAttnPixelDecoder(nn.Module):
# ...
@autocast(enabled=False)
def forward_features(self, features):
srcs = []
pos = []
# Reverse feature maps into top-down order (from low to high resolution)
for idx, f in enumerate(self.transformer_in_features[::-1]):
x = features[f].float() # deformable detr does not support half precision
srcs.append(self.input_proj[idx](x)) # 报错在这一行
pos.append(self.pe_layer(x))
#...
检查这里的features,得到:
1
2
3
4
torch.Size([1, 24, 336, 200])
torch.Size([1, 40, 168, 100])
torch.Size([1, 80, 84, 50])
torch.Size([1, 112, 84, 50])
作为对照,resnet的输出为:
1
2
3
4
torch.Size([1, 256, 160, 160])
torch.Size([1, 512, 80, 80])
torch.Size([1, 1024, 40, 40])
torch.Size([1, 2048, 20, 20])
这里出现一个怪异的问题,我明明设置的out_features=[24, 40, 80, 160],为什么得到的out_features=[24, 40, 80, 112]?
经过逐层的检查,最后还是回到了mobilenetv3,打印了Torch中mobilenetv3-large的各层次。
它共包含有20层。以[channel, scale]来记,层2-3为[16, 1/2],层4-5为[24, 1/4],层6-8为[40, 1/8],层9-12为[80, 1/16],层13-14为[112, 1/16],层15-17为[160, 1/32],层18-20为最后分类的层。
所以,还是太相信AI导致的了。
mobilenetv3没有像resnet那样的res1-res5,而是这样的bottleneck块。交给AI时错误地把112层(13-14)作为了层5。而根据通道翻倍尺寸长宽折半的原则,为了提取到160,要往下额外提取一级。最终枚举出res2-res5应该是层5、8、12、17,而在先前的代码中错误地选为了5、8、12、14。
最终形成的完整可训练架构如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@BACKBONE_REGISTRY.register()
class D2PretrainedMobileNetV3(Backbone):
def __init__(self, cfg, input_shape):
super().__init__()
self._out_features = cfg.MODEL.MOBILENET.OUT_FEATURES
# 加载预训练的 MobileNetV3-Large 模型
self.mobilenet = models.mobilenet_v3_large(pretrained=True)
# 截断模型,只保留到倒数第二层(不包括最后的分类层)
self.features = nn.Sequential(*list(self.mobilenet.features.children()))
# 定义特征层的通道数和步长
self._out_feature_channels = {
"res2": 24,
"res3": 40,
"res4": 80,
"res5": 160,
}
self._out_feature_strides = {
"res2": 4,
"res3": 8,
"res4": 16,
"res5": 32,
}
def forward(self, x):
"""
Args:
x: Tensor of shape (N, C, H, W). H, W must be a multiple of ``self.size_divisibility``.
Returns:
dict[str->Tensor]: names and the corresponding features
"""
assert x.dim() == 4, f"MobileNetV3 takes an input of shape (N, C, H, W). Got {x.shape} instead!"
outputs = {}
for i, layer in enumerate(self.features):
x = layer(x) # 下面的数要+2才会对应上原始的层次
if i == 3: # res2
outputs["res2"] = x
elif i == 6: # res3
outputs["res3"] = x
elif i == 10: # res4
outputs["res4"] = x
elif i == 15: # res5
outputs["res5"] = x
return outputs
def output_shape(self):
return {
name: ShapeSpec(
channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
)
for name in self._out_features
}
@property
def size_divisibility(self):
return 32
这个问题顺利解决。
那你肯定要问了,为什么这么简单一个问题早没看出来。
事实上就是,报错隔得确实特别远,即便在backbone写错,根本不会报在backbone的问题。
加上对数字确实太不敏感了。
再探Criterion
既然整个loss是一个加权相加的结构,那么就意味着,如果设置SUBCAT_WEIGHT=0,从逻辑上说,它应该表现出正常的预测性能,包括main category和mask,只是不能预测出subcategories。
然而,无论如何训练,目前无法得出正常的训练结果,除了能保证工作流外完全无法正常进行训练。除开上述两个阶段的问题,在进行训练,会有一定概率诱发“因为Criterion错误而异常中断”的状况,且在SUBCAT_WEIGHT取较低值(如与原始模型中no_cat相似取0.01,或取0)时相对高发。因此需要进一步检查这一部分的内容。