使用自定义特征提取器计算FID的常见陷阱与解决方案

本文深入探讨了在使用 `torchmetrics` 库计算 fid 时，将自定义 `nn.module` 作为特征提取器可能遇到的数据类型不匹配问题。通过分析 `runtimeerror: expected scalar type byte but found float` 错误，文章阐明了 pytorch 模型通常期望浮点数输入而非字节类型，并提供了将输入图像数据正确转换为浮点数并进行归一化的解决方案及示例代码，旨在帮助开发者顺利集成自定义模型进行 fid 计算。

理解FID与自定义特征提取器

Frechet Inception Distance (FID) 是衡量生成图像质量的重要指标，它通过比较真实图像和生成图像在特征空间中的分布来评估两者的相似性。torchmetrics 库提供了 FrechetInceptionDistance 类，允许用户灵活地使用自定义的 nn.Module 作为特征提取器，而非默认的预训练 Inception-v3 模型。这为研究人员和开发者提供了极大的便利，可以根据特定任务或数据特性定制特征提取逻辑。

当使用自定义特征提取器时，其核心思想是将图像输入到自定义模型中，模型输出一个特征向量，然后 torchmetrics 利用这些特征向量来计算 FID。然而，在集成自定义模型时，一个常见的陷阱是输入数据类型与模型期望不符，从而导致运行时错误。

常见错误：数据类型不匹配

考虑以下使用 torchvision.models.inception_v3 作为自定义特征提取器的示例代码：

import torch
_ = torch.manual_seed(123)
from torchmetrics.image.fid import FrechetInceptionDistance
from torchvision.models import inception_v3, Inception_V3_Weights

# 加载预训练的Inception-v3模型，并设置为评估模式
# 注意：InceptionV3_Weights.DEFAULT 包含了预处理逻辑
net = inception_v3(weights=Inception_V3_Weights.DEFAULT)
net.eval()

# 尝试使用自定义Inception-v3模型初始化FID
fid = FrechetInceptionDistance(feature=net)

# 生成两组模拟图像数据
# 错误示范：使用torch.uint8类型
imgs_dist1 = torch.randint(0, 200, (100, 3, 299, 299), dtype=torch.uint8)
imgs_dist2 = torch.randint(100, 255, (100, 3, 299, 299), dtype=torch.uint8)

fid.update(imgs_dist1, real=True)
fid.update(imgs_dist2, real=False)
result = fid.compute()

print(result)

运行上述代码，会遇到类似以下的 RuntimeError：

RuntimeError: expected scalar type Byte but found Float

这个错误信息清晰地指出，模型期望的输入是 Float 类型，但实际接收到的是 Byte 类型（即 torch.uint8）。

错误根源分析

torchvision 提供的预训练模型，包括 Inception-v3，通常期望的输入图像是 torch.float32 类型，且像素值被归一化到特定范围（例如 [0, 1] 或 [-1, 1]）。这是深度学习模型训练和推理的常见实践。

在上述错误示例中，imgs_dist1 和 imgs_dist2 是通过 torch.randint 生成的，并明确指定了 dtype=torch.uint8。当 FrechetInceptionDistance 内部尝试将这些 uint8 类型的图像数据传递给作为特征提取器的 net（即 Inception-v3 模型）时，Inception-v3 模型在进行卷积操作时发现输入数据类型与期望的 float 类型不匹配，从而抛出 RuntimeError。

绘蛙

电商场景的AI创作平台，无需高薪聘请商拍和文案团队，使用绘蛙即可低成本、批量创作优质的商拍图、种草文案

下载

解决方案：正确的数据类型与预处理

解决这个问题的关键在于确保输入到特征提取器的数据类型和范围与模型训练时所使用的保持一致。对于 torchvision 的模型，这意味着将图像数据转换为 torch.float32 类型，并进行适当的归一化。

以下是修正后的代码示例：

import torch
_ = torch.manual_seed(123)
from torchmetrics.image.fid import FrechetInceptionDistance
from torchvision.models import inception_v3, Inception_V3_Weights
import torchvision.transforms as T

# 加载预训练的Inception-v3模型，并设置为评估模式
# 使用Inception_V3_Weights.DEFAULT加载模型和其默认的预处理权重
weights = Inception_V3_Weights.DEFAULT
net = inception_v3(weights=weights)
net.eval()

# 获取InceptionV3模型期望的预处理变换
# 注意：torchmetrics的FID类在内部对InceptionV3特征提取器有特定的处理，
# 如果feature是标准的InceptionV3模型，它可能会自动处理归一化。
# 但对于自定义模型，用户需要手动确保预处理。
# 对于Inception_V3_Weights.DEFAULT，其预处理通常包括Resize, ToTensor, Normalize。
# 在这里我们主要关注从uint8到float的转换和基本的归一化。

# 初始化FID，如果feature是标准的InceptionV3，torchmetrics会进行内部处理
# 但如果传递的是一个通用的nn.Module，则需要外部确保输入符合模型要求
fid = FrechetInceptionDistance(feature=net)

# 生成两组模拟图像数据 (依然是uint8作为原始数据)
imgs_dist1_uint8 = torch.randint(0, 200, (100, 3, 299, 299), dtype=torch.uint8)
imgs_dist2_uint8 = torch.randint(100, 255, (100, 3, 299, 299), dtype=torch.uint8)

# 关键步骤：将uint8数据转换为float32并归一化到[0, 1]
# 注意：PyTorch模型通常期望像素值在[0, 1]或[-1, 1]之间
imgs_dist1_float = imgs_dist1_uint8.to(torch.float32) / 255.0
imgs_dist2_float = imgs_dist2_uint8.to(torch.float32) / 255.0

# 更新FID计算
fid.update(imgs_dist1_float, real=True)
fid.update(imgs_dist2_float, real=False)
result = fid.compute()

print(f"计算得到的FID值: {result}")

在这个修正后的代码中，我们执行了两个关键操作：

1. 数据类型转换: imgs_dist1_uint8.to(torch.float32) 将张量的数据类型从 torch.uint8 转换为 torch.float32。
2. 像素值归一化: / 255.0 将像素值从 [0, 255] 范围归一化到 [0, 1] 范围。这是 torchvision 预训练模型通常期望的输入范围。

通过这些修改，输入数据现在符合 Inception-v3 模型的期望，错误得以解决，FID 计算可以正常进行。

注意事项与最佳实践

1. 始终检查模型期望的输入: 在使用任何自定义 nn.Module 作为特征提取器时，务必查阅其文档或检查其代码，了解它期望的输入数据类型、维度顺序（例如 NCHW 或 NHWC）以及像素值范围。
2. 数据预处理一致性: 确保用于 FID 计算的图像预处理与特征提取器模型训练时所使用的预处理完全一致。这包括图像大小调整、裁剪、归一化（均值和标准差）等。torchvision.transforms 模块是进行这些操作的有力工具。
3. torchmetrics 对 InceptionV3 的特殊处理: torchmetrics.image.fid.FrechetInceptionDistance 对于 torchvision.models.inception_v3 有一些内置的优化和处理。如果直接将 inception_v3(weights=Inception_V3_Weights.DEFAULT) 传递给 feature 参数，torchmetrics 可能会在内部处理一些预处理步骤。然而，当您使用完全自定义的 nn.Module 时，这些内置的便利性将不再适用，您需要自行管理所有输入预处理。
4. 特征提取器的输出: 确保您的自定义特征提取器输出的是一个扁平化的特征向量，或者 torchmetrics 能够理解并将其转换为扁平特征向量的格式。通常，这指的是模型在分类层之前输出的特征。

总结

在使用 torchmetrics 结合自定义 nn.Module 计算 FID 时，数据类型不匹配是一个常见但容易解决的问题。核心在于理解深度学习模型对输入数据类型和范围的期望。通过将 torch.uint8 类型的原始图像数据正确地转换为 torch.float32 并进行适当的归一化，可以确保特征提取器顺利运行，从而获得准确的 FID 结果。始终遵循模型期望的输入规范，是成功集成自定义模型的关键。