关于Torchmetric SSIM计算结果与batch size相关的疑问

代码示例 #

ssim = StructuralSimilarityIndexMeasure(kernel_size=(5, 5))

A = torch.zeros([8, 1, 500, 500])
B = torch.randn([8, 1, 500, 500])

print(ssim(A, B))

ssim1 = ssim(A[0].unsqueeze(0), B[0].unsqueeze(0))
ssim2 = ssim(A[1].unsqueeze(0), B[1].unsqueeze(0))
ssim3 = ssim(A[2].unsqueeze(0), B[2].unsqueeze(0))
ssim4 = ssim(A[3].unsqueeze(0), B[3].unsqueeze(0))
ssim5 = ssim(A[4].unsqueeze(0), B[4].unsqueeze(0))
ssim6 = ssim(A[5].unsqueeze(0), B[5].unsqueeze(0))
ssim7 = ssim(A[6].unsqueeze(0), B[6].unsqueeze(0))
ssim8 = ssim(A[7].unsqueeze(0), B[7].unsqueeze(0))

print((ssim1 + ssim2 + ssim3 + ssim4 + ssim5 + ssim6 + ssim7 + ssim8) / 8)

输出结果 #

tensor(0.0404)
tensor(0.0340)

环境信息 #

• Python版本：3.8.10
• TorchMetrics版本：0.9.1
• PyTorch版本：1.10.1+cu113

原因分析 #

TorchMetrics中的StructuralSimilarityIndexMeasure是有状态的Metric类，会在多次调用之间累积计算的统计量（比如总SSIM值、样本计数）：

1. 第一次调用ssim(A, B)时，Metric计算出8个样本的SSIM均值（即输出的0.0404），同时将这8个样本的统计数据（总SSIM和样本数8）保存在内部状态中。
2. 后续逐个调用ssim(A[i].unsqueeze(0), B[i].unsqueeze(0))时，Metric不会自动重置状态，而是将当前单个样本的SSIM值追加到之前的累积统计中，返回的是所有已处理样本的均值（比如ssim1是9个样本的均值，ssim2是10个样本的均值，以此类推）。
3. 最后将这些累积均值再取平均，结果自然和第一次的批量均值不一致。

解决方法 #

有两种方式可以得到一致的结果：

方法1：每次计算单个样本时使用新的Metric实例 #

# 批量计算
ssim_batch = StructuralSimilarityIndexMeasure(kernel_size=(5, 5))
print(ssim_batch(A, B))

# 逐个计算并取均值
ssim_sum = 0.0
for i in range(8):
    ssim_single = StructuralSimilarityIndexMeasure(kernel_size=(5, 5))
    ssim_sum += ssim_single(A[i].unsqueeze(0), B[i].unsqueeze(0))
print(ssim_sum / 8)

方法2：每次计算前重置Metric的内部状态 #

ssim = StructuralSimilarityIndexMeasure(kernel_size=(5, 5))

# 批量计算
print(ssim(A, B))
# 重置状态
ssim.reset()

# 逐个计算并取均值
ssim_sum = 0.0
for i in range(8):
    ssim_val = ssim(A[i].unsqueeze(0), B[i].unsqueeze(0))
    ssim_sum += ssim_val
    # 重置状态确保下一次是独立计算
    ssim.reset()
print(ssim_sum / 8)

内容的提问来源于stack exchange，提问作者CaipiDE