当前位置：首页 > news >正文

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

news 2025/6/26 23:55:19

在深度学习和计算机视觉领域，张量的形状转换是一个常见的操作。本文将详细讲解一个用于形状转换的函数 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_shapeassert len(x.shape) == 3B, L, C = x.shapeassert 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]。

详细步骤

获取高度和宽度：
```
H, W = hw_shape
```
从输入参数 hw_shape 中提取高度 H 和宽度 W。
断言输入张量的形状：
```
assert len(x.shape) == 3
```
确保输入张量 x 的形状是 [N, L, C]。
提取输入张量的形状参数：
```
B, L, C = x.shape
```
分别获取批量大小 B、序列长度 L 和通道数 C。
验证序列长度是否匹配：
```
assert L == H * W, 'The seq_len doesn\'t match H, W'
```
确保序列长度 L 等于高度 H 乘以宽度 W。这是必要条件，因为序列长度需要与特征图的总像素数匹配。
转换形状：
```
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 torchx = 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]