pytorch学习笔记

 

1.

if self.stu_preact:
    x = feature_student["preact_feats"] + [
        feature_student["pooled_feat"].unsqueeze(-1).unsqueeze(-1)
    ]
else:
    x = feature_student["feats"] + [
        feature_student["pooled_feat"].unsqueeze(-1).unsqueeze(-1)
    ]

这段代码主要是为了将学生网络的特征拼接起来，并准备输入到注意力块中进行处理。

其中，输入 feature_student 是学生网络的特征，

包含了学生网络最后一层卷积的特征以及经过全局池化后的特征。

如果 self.stu_preact 为 True，表示使用学生网络中间的特征，这些特征通常更加丰富，具有更高的表达能力。

因此，x 的第一部分是 feature_student["preact_feats"]，即学生网络中间的特征，第二部分是 feature_student["pooled_feat"] 经过 unsqueeze(-1) 和 unsqueeze(-1) 后的结果，将其变成 4D 张量。

这样，x 就是一个列表，包含了两个元素，分别是 4D 张量和 3D 张量。

这里的 unsqueeze(-1) 表示在张量最后添加一个维度，而这个维度的大小是 1。

由于 unsqueeze() 的默认参数是 dim=0，表示在最前面添加一个维度，所以这里需要指定 dim=-1。因此，feature_student["pooled_feat"].unsqueeze(-1).unsqueeze(-1) 表示将 feature_student["pooled_feat"] 变成 4D 张量，且最后两个维度的大小都是 1。

最后，如果 self.stu_preact 为 True，则将 feature_student["preact_feats"] 和 feature_student["pooled_feat"] 的 4D 张量拼接在一起，

否则将 feature_student["feats"] 和 feature_student["pooled_feat"] 的 4D 张量拼接在一起。

2.

out_features, res_features = self.abfs[0](x[0], out_shape=self.out_shapes[0])
results.append(out_features)

for features, abf, shape, out_shape in zip(
    x[1:], self.abfs[1:], self.shapes[1:], self.out_shapes[1:]
):
    out_features, res_features = abf(features, res_features, shape, out_shape)
    results.insert(0, out_features)
feature_teacher = feature_teacher["preact_feats"][1:] + [
    feature_teacher["pooled_feat"].unsqueeze(-1).unsqueeze(-1)
]

这段代码是ReviewKD的前向传播过程中，

将feature maps输入到注意力机制中，得到经过注意力机制处理后的输出结果。

首先，将学生网络的feature maps输入到第一个注意力机制中（即self.abfs[0]），

得到输出结果out_features和res_features，

其中out_features是经过注意力机制处理后的feature maps，

        res_features是注意力机制中的残差特征。

然后，将处理过的输出结果out_features添加到results列表中。

接着，对于剩下的feature maps，依次输入到对应的注意力机制中（即self.abfs[1:]），

得到经过注意力机制处理后的输出结果out_features和res_features，

并将out_features添加到results的开头。

最后，将教师网络的feature maps（即feature_teacher）中的preact_feats特征与pooled_feat特征进行拼接，得到x列表。

这里注意，由于在上面的注意力机制处理中，将处理后的输出结果out_features添加到了results列表中，

所以在循环处理x[1:]时需要使用insert方法将out_features添加到results的开头。

关键的部分好难懂啊 加油

self.abfs[0](x[0], out_shape=self.out_shapes[0])详细解释一下，

写个方法的demo

这里是使用Attention Block模块，对输入特征x[0]进行特征提取和通道压缩，将其压缩成输出形状为self.out_shapes[0]的特征，同时还输出压缩过程中的残差特征res_features，以备后续的Attention Block模块使用。具体的实现细节可以参考Attention Block模块的具体实现。

以下是一个简单的演示Attention Block模块的示例代码：

import torch
import torch.nn as nn

class AttentionBlock(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels):
        super(AttentionBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(mid_channels, out_channels, kernel_size=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, res, shape, out_shape):
        y = self.conv1(x)
        y = self.bn1(y)
        y = self.relu1(y)
        y = self.conv2(y)
        y = self.bn2(y)
        y = self.sigmoid(y)
        y = nn.functional.interpolate(y, size=shape, mode="bilinear", align_corners=True)
        out = x * y + res
        out = nn.functional.interpolate(out, size=out_shape, mode="bilinear", align_corners=True)
        out = self.conv1(out)
        out = self.bn1(out)
        out = self.relu1(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.sigmoid(out)
        out = nn.functional.interpolate(out, size=out_shape, mode="bilinear", align_corners=True)
        return out, x * out

in_channels = 64
mid_channels = 32
out_channels = 32
input_shape = (4, 64, 64)

x = torch.randn(input_shape)
abf = AttentionBlock(in_channels, mid_channels, out_channels)
out_shape = (4, 32, 32)

out_features, res_features = abf(x, None, input_shape[1:], out_shape)
print(out_features.shape, res_features.shape)  # (4, 32, 32) (4, 64, 64)

这个代码看看就好 ，能懂个大概，因为张量问题直接报错

for features, abf, shape, out_shape in zip(
    x[1:], self.abfs[1:], self.shapes[1:], self.out_shapes[1:]
):

这段代码中使用了Python内置的zip()函数，它会将多个可迭代对象打包成一个元组序列，返回一个zip对象。

在这个例子中，x[1:]是一个列表，包含了输入特征和池化后的特征，self.abfs[1:]是一个包含了多个注意力模块的列表，

self.shapes[1:]是一个输入特征形状的列表，

self.out_shapes[1:]是一个输出特征形状的列表。

zip()函数会将它们一一对应地打包成一个元组，然后返回一个包含多个元组的zip对象。

在for循环中，我们使用了元组拆包的方式，

将每个元组的四个元素分别赋值给了features、abf、shape和out_shape。

然后我们就可以调用abf的forward()方法来对输入特征进行处理，

这个方法会返回处理后的输出特征和残差特征，

我们将输出特征插入到results列表的开头。

接着，我们将残差特征保存到res_features中，用于后续处理。

最后，我们将feature_teacher也按照和x一样的方式进行处理，

得到它的输入特征和池化后的特征，用于计算distillation loss。

feature_teacher = feature_teacher["preact_feats"][1:] + [
    feature_teacher["pooled_feat"].unsqueeze(-1).unsqueeze(-1)

这行代码的作用是将teacher模型提取到的特征按照和student模型类似的方式进行处理。

如果self.stu_preact为True，即student模型最后一层是pre-activation模式，

那么就将teacher模型的pre-activation feature和pooled feature进行拼接，类似于student模型中的x；

如果self.stu_preact为False，即student模型最后一层是post-activation模式，那么就将teacher模型的activation feature和pooled feature进行拼接。

在拼接完成后，teacher模型的特征被处理成了一个类似于student模型中的x的列表，之后将使用attention block对这个列表中的特征进行处理。

feature_teacher["preact_feats"][1:]这个[1:]什么意思

feature_teacher["preact_feats"] 是一个列表，

[1:] 表示取列表的第二个元素开始到最后一个元素（不包括第一个元素）。

所以feature_teacher["preact_feats"][1:] 返回的是feature_teacher["preact_feats"] 列表中除第一个元素以外的元素。

loss_reviewkd = (
    self.reviewkd_loss_weight
    * min(kwargs["epoch"] / self.warmup_epochs, 1.0)
    * hcl_loss(results, feature_teacher)
)

这段代码是在计算 ReviewKD 的损失函数，

其中使用了一个 warm-up 的策略，即在前几个 epoch 中逐渐增加 ReviewKD 损失的权重。

这个权重逐渐增加的方式是使用了 min(kwargs["epoch"] / self.warmup_epochs, 1.0) 这个表达式，

其中 kwargs["epoch"] 是当前训练的 epoch 数，

self.warmup_epochs 是指定的 warm-up 的 epoch 数。

这个表达式会返回一个值，这个值的范围是 0 到 1，

表示当前 epoch 的权重比例，越接近 1，ReviewKD 的损失在总损失中的占比就越大。

在计算损失时，使用了一个名为 hcl_loss 的函数来计算 ReviewKD 的损失，

其中 results 是当前模型的输出结果，feature_teacher 是教师模型的输出结果。