转换张量形状：`nlc_to_nchw` 函数详解

在深度学习和计算机视觉领域，张量的形状转换是一个常见的操作。本文将详细讲解一个用于形状转换的函数 nlc_to_nchw，它能够将形状为 [N, L, C] 的张量转换为 [N, C, H, W] 的张量。

函数定义

def nlc_to_nchw(x, hw_shape):
    """Convert [N, L, C] shape tensor to [N, C, H, W] shape tensor.

    Args:
        x (Tensor): The input tensor of shape [N, L, C] before conversion.
        hw_shape (Sequence[int]): The height and width of output feature map.

    Returns:
        Tensor: The output tensor of shape [N, C, H, W] after conversion.
    """
    H, W = hw_shape
    assert len(x.shape) == 3
    B, L, C = x.shape
    assert L == H * W, 'The seq_len doesn\'t match H, W'
    return x.transpose(1, 2).reshape(B, C, H, W)

参数解释

• x (Tensor): 输入张量，其形状为 [N, L, C]。其中：

  • N 表示批量大小（batch size）。
  • L 表示序列长度（sequence length）。
  • C 表示通道数（channels）。

• hw_shape (Sequence[int]): 输出特征图的高度和宽度，为一个包含两个整数的序列（或元组），即 [H, W]。其中：

  • H 表示高度（height）。
  • W 表示宽度（width）。

返回值

• Tensor: 转换后的张量，形状为 [N, C, H, W]。

详细步骤

1. 获取高度和宽度：

   H, W = hw_shape
   

   从输入参数 hw_shape 中提取高度 H 和宽度 W。

2. 断言输入张量的形状：

   assert len(x.shape) == 3
   

   确保输入张量 x 的形状是 [N, L, C]。

3. 提取输入张量的形状参数：

   B, L, C = x.shape
   

   分别获取批量大小 B、序列长度 L 和通道数 C。

4. 验证序列长度是否匹配：

   assert L == H * W, 'The seq_len doesn\'t match H, W'
   

   确保序列长度 L 等于高度 H 乘以宽度 W。这是必要条件，因为序列长度需要与特征图的总像素数匹配。

5. 转换形状：

   return x.transpose(1, 2).reshape(B, C, H, W)
   

   通过 transpose 和 reshape 操作，将输入张量从 [N, L, C] 转换为 [N, C, H, W]：

   • x.transpose(1, 2) 将张量的第二个和第三个维度交换，形状变为 [N, C, L]。
   • reshape(B, C, H, W) 将张量重新调整形状为 [N, C, H, W]。

示例

假设有一个形状为 [2, 6, 3] 的输入张量 x，其中 N=2, L=6, C=3。给定的 hw_shape 为 [2, 3]，表示高度 H=2 和宽度 W=3。

import torch

x = torch.rand(2, 6, 3)  # 创建一个形状为 [2, 6, 3] 的随机张量
hw_shape = [2, 3]        # 定义高度和宽度

output = nlc_to_nchw(x, hw_shape)  # 调用函数
print(output.shape)  # 输出张量的形状为 [2, 3, 2, 3]

运行上述代码后，输出张量的形状将为 [2, 3, 2, 3]，符合预期的 [N, C, H, W] 形状。

总结

通过 nlc_to_nchw 函数，我们可以轻松地将形状为 [N, L, C] 的张量转换为 [N, C, H, W] 的张量。这在处理图像数据和构建神经网络时非常有用，因为不同的层和操作可能要求特定的张量形状。理解和掌握这些基本的张量操作是深度学习实践中的重要技能。