解析PyTorch代码(img.sum(dim=1)!=0).float()[:, None, :, :]作用

目标代码如下：

(img.sum(dim=1)!=0).float()[:, None, :, :]

这段代码来自ViLT的视觉编码器模块，作用是生成图像的有效区域掩码，逐部分拆解如下：

• 前置前提：代码中输入的img是预处理后的4维图像张量，标准形状为[batch_size, 3, image_height, image_width]，四个维度分别对应批次大小、RGB三颜色通道、图像高度、图像宽度，图像补边的填充区域像素值全为0。
• img.sum(dim=1)：你对sum操作的作用判断大方向没错，但dim=1对应的维度搞混了。这里dim=1是通道维度，不是水平列维度，沿该维度求和就是把每个像素位置的R、G、B三个通道值相加，输出形状为[batch_size, image_height, image_width]的3维张量，每个位置的值对应该像素三通道的数值总和。填充区域因为三个通道全是0，求和结果自然为0，真实图像区域的像素求和结果几乎不可能为0。
• !=0：你的理解完全正确。这是张量逐元素比较操作，对求和结果的每个位置做非零判断，非零位置返回布尔值True（对应真实图像内容），零值位置返回False（对应填充区域），输出形状保持[batch_size, image_height, image_width]不变，类型为布尔张量。
• .float()：你的理解也正确。作用是把布尔张量转为浮点型张量，True映射为1.0，False映射为0.0，形状不变，主要是为了和后续网络中浮点类型的张量做运算兼容，避免类型报错。
• [:, None, :, :]索引操作：你觉得前面输出维度低是误解，前面三步输出的是3维张量，索引里的三个:分别对应原3维张量的三个维度，意思是这三个维度的所有元素全部保留，不做任何筛选。
  索引中写None是PyTorch继承自NumPy的语法糖，作用是在None所在的位置插入一个长度为1的新维度。这里None放在第二个索引位，也就是在批次维度之后、高度维度之前插入长度为1的维度，最终整个代码段输出的张量形状为[batch_size, 1, image_height, image_width]。
  插入这个维度的目的是满足后续广播运算的要求：后续视觉Transformer中做注意力掩码、patch特征加权时，用到的张量都是4维结构，插入长度为1的维度后，这个掩码张量不需要额外reshape就能自动广播对齐目标张量的形状，简化代码写法。

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