【小白入门】超详细的OCRnet详解(含代码分析)
本文仅梳理总结自己在学习过程中的一些理解和思路,不保证绝对正确,请酌情参考。如果各位朋友发现任何错误请及时告诉我,大家一起讨论共同提高。
本文参考博客https://blog.csdn.net/u011622208/article/details/110202659 及http://yearing1017.cn/2020/09/07/OCRNet/,有部分内容为直接引用,如有侵权,烦请告知(我删的贼快)
论文:https://arxiv.org/pdf/1909.11065.pdf
已开源代码地址:https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/HRNet-OCR
OCRnet
针对语义分割中如何构建上下文信息,微软亚洲研究院和中科院计算所的研究员们提出了一种新的物体上下文信息——在构建上下文信息时显式地增强了来自于同一类物体的像素的贡献,这种新的上下文信息从语义分割的定义出发,符合第一性原理思维,在2019年7月和2020年1月的 Cityscapes leaderboard 提交结果中都取得了语义分割任务第一名的成绩。相关工作“Object-Contextual Representations for Semantic Segmentation”已经被 ECCV 2020 收录。
简介
对于语义分割任务来说,其两大关键是:分辨率和上下文。
- 语义分割一个密集像素预测任务,因此空间分辨率很重要。
- 像素本身不具备语义,它的语义由其图像整体或目标区域决定,因此它对上下文高度依赖。
- 一个像素位置的上下文指的是它周围的像素位置。
该论文的主要思想也就是像素的类别标签是由它所在的目标的类别标签决定的。主要思路是利用目标区域表示来增强其像素的表示。与之前的考虑上下文关系的方法不同的是,之前的方法考虑的是上下文像素之间的关系,没有显示利用目标区域的特征。
网络结构
官方给出的网络结构图如图
其中,粉红色虚线框内为形成的软对象区域(Soft Object Regions),紫色虚线框中为物体区域表示(Object Region Representations),橙色虚线框中为对象上下文表示和增强表示。
具体实现(含代码分析)
论文中指出,OCR 方法的实现主要包括3个阶段:
First, we divide the contextual pixels into a set of soft object regions with each corresponding to a class, i.e., a coarse soft segmentation computed from a deep network (e.g., ResNet [23] or HRNet [55]). Such division is learned under the supervision of the ground-truth segmentation.
第一步: 将上下文像素划分为一组软对象区域,每个soft object regions对应一个类,即从深度网络(backbone)计算得到的粗软分割(粗略的语义分割结果)。这种划分是在ground-truth分割的监督下学习的。根据网络中间层的特征表示估测粗略的语义分割结果作为 OCR 方法的一个输入,即结构图中粉红色框内的Soft Object Regions
self.aux_head = nn.Sequential(
nn.Conv2d(high_level_ch, high_level_ch,
kernel_size=1, stride=1, padding=0),
BNReLU(high_level_ch),
nn.Conv2d(high_level_ch, num_classes,
kernel_size=1, stride=1, padding=0, bias=True)
)
aux_out = self.aux_head(high_level_features) #soft object regions
#high_level_features为backbone输出的粗略的高层特征结果
将backbone的输出结果经过1*1的卷积后输出b×k×h×w的张量作为软对象区域,其中,k为粗略分类后对象的类别数(eg:若有17个类别,则网络输出为:b×17×h×w)
Second, we estimate the representation for each object region by aggregating the representations of the pixels in the corresponding object region.
第二步: 根据粗略的语义分割结果(soft object regions)和网络最深层输出的像素特征(Pixel Representations)表示计算出 K 组向量,即物体区域表示(Object Region Representations),其中每一个向量对应一个语义类别的特征表示
经评论区@duoheshuiya,@qq_39680835两位大佬指正,此步aux-out出来的c应该是类别数,k为预设的通道数。非常感谢两位的纠正,也向修改之前被我误导各位表示歉意
def forward(self, feats, probs):
batch_size, c, _, _ = probs.size(0), probs.size(1), probs.size(2), \
probs.size(3)
# each class image now a vector
probs = probs.view(batch_size, c, -1)
feats = feats.view(batch_size, feats.size(1), -1)
feats = feats.permute(0, 2, 1) # batch x hw x c
probs = F.softmax(self.scale * probs, dim=2) # batch x k x hw
ocr_context = torch.matmul(probs, feats)
ocr_context = ocr_context.permute(0, 2, 1).unsqueeze(3)
return ocr_context
其中,feats为网络最深层输出的像素特征,由backbone输出的粗略的高层特征结果经卷积处理得到。在调用此函数时,将aux-out作为实参传给probs。
self.conv3x3_ocr = nn.Sequential(
nn.Conv2d(high_level_ch, ocr_mid_channels,
kernel_size=3, stride=1, padding=1),
BNReLU(ocr_mid_channels),
)
feats = self.conv3x3_ocr(high_level_features)
这段代码的功能如图
经上述处理,得到bck的张量ocr-context,其中,b为batch-size,c为类别个数,每个类别可用一个大小为c的向量来描述。这个描述则为下文的注意力机制提供了参照信息(query)
Last, we augment the representation of each pixel with the object-contextual representation (OCR). The OCR is the weighted aggregation of all the object region representations with the weights calculated according to the relations between pixels and object regions.
第三步: 这一步可以再细分为两个步骤
- 计算网络最深层输出的像素特征表示(Pixel Representations)与计算得到的物体区域特征表示(Object Region Representation)之间的关系矩阵,然后根据每个像素和物体区域特征表示在关系矩阵中的数值把物体区域特征加权求和,得到最后的物体上下文特征表示 ** (Object Contextual Representation),即OCR** 。
class ObjectAttentionBlock(nn.Module):
'''
The basic implementation for object context block
Input:
N X C X H X W
Parameters:
in_channels : the dimension of the input feature map
key_channels : the dimension after the key/query transform
scale : choose the scale to downsample the input feature
maps (save memory cost)
Return:
N X C X H X W
'''
def __init__(self, in_channels, key_channels, scale=1):
super(ObjectAttentionBlock, self).__init__()
self.scale = scale
self.in_channels = in_channels
self.key_channels = key_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_pixel = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
BNReLU(self.key_channels),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
BNReLU(self.key_channels),
)
self.f_object = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
BNReLU(self.key_channels),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
BNReLU(self.key_channels),
)
self.f_down = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0, bias=False),
BNReLU(self.key_channels),
)
self.f_up = nn.Sequential(
nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels,
kernel_size=1, stride=1, padding=0, bias=False),
BNReLU(self.in_channels),
)
def forward(self, x, proxy):
batch_size, h, w = x.size(0), x.size(2), x.size(3)
if self.scale > 1:
x = self.pool(x)
query = self.f_pixel(x).view(batch_size, self.key_channels, -1)
query = query.permute(0, 2, 1)
key = self.f_object(proxy).view(batch_size, self.key_channels, -1)
value = self.f_down(proxy).view(batch_size, self.key_channels, -1)
value = value.permute(0, 2, 1)
sim_map = torch.matmul(query, key)
sim_map = (self.key_channels**-.5) * sim_map
sim_map = F.softmax(sim_map, dim=-1)
# add bg context ...
context = torch.matmul(sim_map, value)
context = context.permute(0, 2, 1).contiguous()
context = context.view(batch_size, self.key_channels, *x.size()[2:])
context = self.f_up(context)
if self.scale > 1:
context = F.interpolate(input=context, size=(h, w), mode='bilinear',
align_corners=cfg.MODEL.ALIGN_CORNERS)
return context
这段代码主要功能如图
通过对网络最深层输出的像素特征表示(Pixel Representations)与物体区域特征表示(Object Region Representation)进行卷积处理,获得query,key和value值,按照下述公式计算相似度得分,其中,dk取dk=key_channels。
这一步体现了注意力机制的应用,也是ocrnet性能提升的关键步骤
- 当把物体上下文特征表示 OCR 与网络最深层输入的特征表示拼接之后作为上下文信息增强的特征表示(Augmented Representation),可以基于增强后的特征表示预测每个像素的语义类别。
class SpatialOCR_Module(nn.Module):
"""
Implementation of the OCR module:
We aggregate the global object representation to update the representation
for each pixel.
"""
def __init__(self, in_channels, key_channels, out_channels, scale=1,
dropout=0.1):
super(SpatialOCR_Module, self).__init__()
self.object_context_block = ObjectAttentionBlock(in_channels,
key_channels,
scale)
if cfg.MODEL.OCR_ASPP:
self.aspp, aspp_out_ch = get_aspp(
in_channels, bottleneck_ch=cfg.MODEL.ASPP_BOT_CH,
output_stride=8)
_in_channels = 2 * in_channels + aspp_out_ch
else:
_in_channels = 2 * in_channels
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0,
bias=False),
BNReLU(out_channels),
nn.Dropout2d(dropout)
)
def forward(self, feats, proxy_feats):
context = self.object_context_block(feats, proxy_feats)
if cfg.MODEL.OCR_ASPP:
aspp = self.aspp(feats)
output = self.conv_bn_dropout(torch.cat([context, aspp, feats], 1))
else:
output = self.conv_bn_dropout(torch.cat([context, feats], 1))
return output
将第三步第一小步中得到的context(物体上下文特征表示 OCR)与feats(网络最深层输出的像素特征)直接用cat拼接,得到上下文信息增强的特征表示。
综上,整个代码结构如下图所示
其中,代码中faets对应网络结构图中的网络最深层输出的像素特征表示(Pixel Representations),aut_out对应粗略的语义分割结果(soft object regions),ocr-context对应物体区域表示(Object Region Representations)
实验结果
OCR 方法在5个主流的语义分割数据库上都取得了不错的结果。值得注意的是,基于 HRNet + OCR 的一个单模型结果在 ECCV 2020 提交截稿前,在 Cityscapse leaderboard 的语义分割任务中取得了排名第一。
综上,OCR 方法提出的物体上下文信息的目的在于显式地增强物体信息,通过计算一组物体的区域特征表达,根据物体区域特征表示与像素特征表示之间的相似度将这些物体区域特征表示传播给每一个像素。在街景分割任务中,OCR 方法也比 PSPNet 的 PPM 和DeepLabv3 的 ASPP更加高效也更加准确。截止到2020年8月20日,根据最新 Cityscapes leaderboard 的结果,来自于 NVIDIA 的研究团队采用 HRNet + OCR 作为主干网络结构并且设计了一种更高效的多尺度融合方法[16],取得了目前排名第一的结果:85.4%。另外在最新的 ADE20K leaderboard 上,来自创新奇智(AInnovation)的研究团队[17]也基于 HRNet + OCR 取得了目前第一名的结果:59.48%。
最后,错误之处,希望各位大佬不吝赐教