PyTorch实现ResNet18量化感知训练推理时遇NotImplementedError

我正尝试基于PyTorch官方文档实现ResNet18的量化感知训练（QAT），但在推理阶段调用evaluate_model函数时触发错误，具体报错如下：

NotImplementedError: Could not run 'aten::add.out' with arguments from the 'QuantizedCPU' backend

完整错误信息：

NotImplementedError: Could not run 'aten::add.out' with arguments from the 'QuantizedCPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::add.out' is only available for these backends: [CPU, CUDA, Meta, MkldnnCPU, SparseCPU, SparseCUDA, SparseMeta, SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMTIA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradMeta, AutogradNestedTensor, Tracer, AutocastCPU, AutocastXPU, AutocastMPS, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot,

完整实现代码 #

Block 1 - 导入必要库、定义训练与评估函数 #

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision.models import resnet18
import matplotlib.pyplot as plt
import copy
import numpy as np
import os

def evaluate_model(model, test_loader, device, criterion=None):
    model.eval()
    model.to(device)
    running_loss = 0
    running_corrects = 0
    for inputs, labels in test_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        if criterion is not None:
            loss = criterion(outputs, labels).item()
        else:
            loss = 0
        running_loss += loss * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
    eval_loss = running_loss / len(test_loader.dataset)
    eval_accuracy = running_corrects / len(test_loader.dataset)
    return eval_loss, eval_accuracy


def train_model(model, train_loader, test_loader, device, learning_rate=1e-1, num_epochs=200):
    criterion = nn.CrossEntropyLoss()
    model.to(device)
    #optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=1e-4)
    optimizer = optim.Adam(model.parameters(), lr=1e-4)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1, last_epoch=-1)
    model.eval()
    eval_loss, eval_accuracy = evaluate_model(model=model, test_loader=test_loader, device=device, criterion=criterion)
    print("Epoch: {:02d} Eval Loss: {:.3f} Eval Acc: {:.3f}".format(-1, eval_loss, eval_accuracy))
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0
        running_corrects = 0
        for inputs, labels in train_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
            if torch.isnan(loss):
                print("NaN in Loss!")
                return model
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)

        train_loss = running_loss / len(train_loader.dataset)
        train_accuracy = running_corrects / len(train_loader.dataset)
        model.eval()
        eval_loss, eval_accuracy = evaluate_model(model=model, test_loader=test_loader, device=device, criterion=criterion)
        scheduler.step()
        print("Epoch: {:03d} Train Loss: {:.3f} Train Acc: {:.3f} Eval Loss: {:.3f} Eval Acc: {:.3f}".format(epoch, train_loss, train_accuracy, eval_loss, eval_accuracy))
    return model

Block 2 - 加载并预处理CIFAR-100数据集（调整为224*224） #

transform_train = transforms.Compose([
    transforms.Resize(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

trainset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False, num_workers=2)

print("Data loaded and transformed successfully!")

Block 3 - 量化模型定义与训练推理 #

class QuantizedResNet18(nn.Module):
    def __init__(self, model_fp32):
        super().__init__()
        # QuantStub converts tensors from floating point to quantized.
        # This will only be used for inputs.
        self.quant = torch.ao.quantization.QuantStub()
        # DeQuantStub converts tensors from quantized to floating point.
        # This will only be used for outputs.
        self.dequant = torch.ao.quantization.DeQuantStub()
        # FP32 model
        self.model_fp32 = model_fp32

    def forward(self, x):
        # manually specify where tensors will be converted from floating
        # point to quantized in the quantized model
        print(f"Input shape before quant: {x.shape}, dtype: {x.dtype}")
        x = self.quant(x)
        print(f"Input shape after quant: {x.shape}, dtype: {x.dtype}")
        x = self.model_fp32(x)
        print(f"Input shape: {x.shape}, dtype: {x.dtype}")
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
        print(f"Input shape: {x.shape}, dtype: {x.dtype}")
        return x

model = resnet18(num_classes=100, pretrained=False)
fused_model = copy.deepcopy(model)
fused_model.eval()
qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
fused_model.qconfig = qconfig

# Fuse the model in place rather manually.
fused_model = torch.ao.quantization.fuse_modules(fused_model, [["conv1", "bn1", "relu"]], inplace=True)
for module_name, module in fused_model.named_children():
    if "layer" in module_name:
        for basic_block_name, basic_block in module.named_children():
            torch.ao.quantization.fuse_modules(basic_block, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True)
            for sub_block_name, sub_block in basic_block.named_children():
                if sub_block_name == "downsample":
                    torch.ao.quantization.fuse_modules(sub_block, [["0", "1"]], inplace=True)

quantized_model_1 = QuantizedResNet18(model_fp32=fused_model)
quantized_model_1.qconfig = qconfig
cuda_device = torch.device("cuda:0")
quantized_model_1_prepared = torch.ao.quantization.prepare_qat(quantized_model_1.train())
trained_quantized_model_1_prepared = train_model(model=quantized_model_1_prepared, train_loader=trainloader, test_loader=testloader, device=cuda_device, learning_rate=1e-3, num_epochs=1)
cpu_device = torch.device("cpu:0")
trained_quantized_model_1_prepared.to(cpu_device)
trained_quantized_model_1_prepared.eval()
trained_quantized_model_1_prepared_int8 =  torch.ao.quantization.convert(trained_quantized_model_1_prepared)
print(evaluate_model(model=trained_quantized_model_1_prepared_int8, test_loader=testloader, device=cpu_device))

问题原因与解决方法 #

这个错误的核心原因是：量化后的模型中存在浮点张量与量化张量直接相加的操作，而QuantizedCPU后端不支持这种混合精度的add运算。在ResNet的残差连接中，当存在下采样分支时，主分支的输出（量化张量）和下采样分支的输出可能未被正确量化，导致类型不匹配。

具体修复步骤：

1. 确保所有子模块都正确设置qconfig
   手动设置顶层模型qconfig容易遗漏子模块，递归给所有子模块配置qconfig：

   # 替换原fused_model.qconfig = qconfig这一行
   def set_qconfig_recursive(model, qconfig):
       model.qconfig = qconfig
       for child in model.children():
           set_qconfig_recursive(child, qconfig)
   set_qconfig_recursive(fused_model, qconfig)
   

2. 修改评估函数，避免破坏量化模型结构
   量化模型（int8）只能在CPU运行，调用model.to(device)会破坏量化张量的结构，调整evaluate_model：

   def evaluate_model(model, test_loader, device, criterion=None):
       model.eval()
       # 量化模型跳过to(device)操作
       if not hasattr(model, 'quant'):
           model.to(device)
       running_loss = 0
       running_corrects = 0
       for inputs, labels in test_loader:
           inputs = inputs.to(device)
           labels = labels.to(device)
           outputs = model(inputs)
           _, preds = torch.max(outputs, 1)
           if criterion is not None:
               loss = criterion(outputs, labels).item()
           else:
               loss = 0
           running_loss += loss * inputs.size(0)
           running_corrects += torch.sum(preds == labels.data)
       eval_loss = running_loss / len(test_loader.dataset)
       eval_accuracy = running_corrects / len(test_loader.dataset)
       return eval_loss, eval_accuracy
   

3. 改用FX模式量化（更可靠）
   手动模块融合容易遗漏分支，FX模式能自动遍历所有子模块处理量化：

   # 替换Block3中的量化相关代码
   model = resnet18(num_classes=100, pretrained=False)
   qconfig = torch.ao.quantization.get_default_qat_qconfig('fbgemm')
   qconfig_dict = {"": qconfig}
   
   # FX模式准备QAT模型
   prepared_model = torch.ao.quantization.prepare_qat_fx(model, qconfig_dict)
   
   # 训练QAT模型（训练可在GPU，转换前需回到CPU）
   trained_model = train_model(model=prepared_model.train(), train_loader=trainloader, test_loader=testloader, device=cuda_device, learning_rate=1e-3, num_epochs=1)
   
   # 转换为量化模型
   trained_model.to(cpu_device)
   trained_model.eval()
   quantized_model = torch.ao.quantization.convert_fx(trained_model)
   
   # 评估
   print(evaluate_model(model=quantized_model, test_loader=testloader, device=cpu_device))
   

关键注意点：

• QAT训练阶段可在GPU运行，但转换为int8量化模型后必须在CPU上推理，量化算子仅支持CPU后端。
• 手动修改模型结构时容易遗漏残差连接等分支的量化配置，FX模式能减少人为错误。

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