class Downsample(nn.Module): """ Down-sampling block based on: "Hatamizadeh et al., FasterViT: Fast Vision Transformers with Hierarchical Attention """ def __init__(self, dim, keep_dim=False, ): """ Args: dim: feature size dimension. norm_layer: normalization layer. keep_dim: bool argument for maintaining the resolution. """ super().__init__() if keep_dim: dim_out = dim else: dim_out = 2 * dim self.norm = LayerNorm2d(dim) self.reduction = nn.Sequential( nn.Conv2d(dim, dim_out, 3, 2, 1, bias=False), ) def forward(self, x): x = self.norm(x) x = self.reduction(x) return x 逐行解释

这是一个名为`Downsample`的PyTorch模块，用于进行下采样操作。它基于论文"FasterViT: Fast Vision Transformers with Hierarchical Attention"中的方法。

在`__init__`方法中，它接受两个参数：`dim`表示特征大小的维度，`keep_dim`是一个布尔值，用于控制是否保持分辨率。

在初始化方法中，它首先使用`LayerNorm2d`对输入进行归一化。然后，它使用一个包含单个卷积层的`reduction`序列来进行下采样操作。这个卷积层的输入通道数为`dim`，输出通道数为`dim_out`，卷积核大小为3x3，步幅为2，填充为1。如果`keep_dim`为True，则输出通道数与输入通道数相同，否则输出通道数为输入通道数的两倍。

在`forward`方法中，它首先对输入进行归一化处理，然后将归一化后的输入传递给下采样层，并返回下采样后的结果。