pytorch 计算二维图像的平均值

在 PyTorch 中，对于一个二维的图像，可以使用 `torch.mean()` 函数来计算平均值。假设你有一个形状为 `(H, W)` 的二维图像张量 `img`，可以使用以下代码计算其平均值：

mean = torch.mean(img)

如果你有一个形状为 `(C, H, W)` 的三维图像张量 `img`，其中 `C` 是通道数，可以使用 `torch.mean()` 函数和 `dim` 参数来计算通道维度上的平均值。例如，要计算每个通道的平均值，可以使用以下代码：

mean = torch.mean(img, dim=(1, 2))

这将返回一个形状为 `(C,)` 的张量，其中每个元素是对应通道的平均值。如果你想计算所有元素的平均值，可以将 `dim` 参数设置为 `None`：

mean = torch.mean(img, dim=None)

这将返回一个标量张量，即所有元素的平均值。注意，在计算平均值之前，你可能需要将图像的数据类型转换为浮点型，以便得到正确的结果。可以使用 `.float()` 方法将张量转换为浮点类型：

img = img.float()

完整的示例代码如下：

import torch

# 假设你有一个形状为 (H, W) 的二维图像张量 img
mean = torch.mean(img)

# 假设你有一个形状为 (C, H, W) 的三维图像张量 img
mean = torch.mean(img, dim=(1, 2))