GPU环境生成Adversarial Patch遇张量设备不匹配RuntimeError求助

报错信息 #

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument tensors in method wrapper_cat)

代码在CPU环境可正常运行，GPU服务器执行时触发上述错误，以下是最小复现代码：

from art.estimators.object_detection.pytorch_yolo import PyTorchYolo
from art.attacks.evasion import AdversarialPatchPyTorch
from inria_utils import load_inria
from evaluation_metrics import evaluate_patch
import torch
from yolov5.utils.loss import ComputeLoss
import yolov5

def load_model():
    class Yolo(torch.nn.Module):
        def __init__(self, model):
            super().__init__()
            self.model = model
            self.model.hyp = {'box': 0.05,
                                'obj': 1.0,
                                'cls': 0.5,
                                'anchor_t': 4.0,
                                'cls_pw': 1.0,
                                'obj_pw': 1.0,
                                'fl_gamma': 0.0
                                }
            self.compute_loss = ComputeLoss(self.model.model.model)

        def forward(self, x, targets=None):
            if self.training:
                outputs = self.model.model.model(x)
                loss, loss_items = self.compute_loss(outputs, targets)
                loss_components_dict = {"loss_total": loss}
                return loss_components_dict
            else:
                return self.model(x)

    # Set the device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    model = yolov5.load('yolov5s.pt')

        
    model = Yolo(model)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return  PyTorchYolo(model=model,
                        device_type= 'cuda' if torch.cuda.is_available() else 'cpu',
                        input_shape=(3, 640, 640),
                        clip_values=(0, 255), 
                        attack_losses=("loss_total",))

def main():

    detector = load_model()

    x, _ = load_inria(subset="train", num_samples=8)#x is a numpy array
    # I have tried x to tensor and set the device to cuda but ap.generate expects np array
    target = detector.predict(x) #its a list

    ap = AdversarialPatchPyTorch(
        estimator=detector,
        rotation_max=8,
        scale_min=0.4,
        scale_max=1,
        learning_rate=1,
        batch_size=16,
        max_iter=5,
        patch_shape=(3, 200, 200),
        patch_type='square',
        verbose=True,
        optimizer='Adam')
    
    ap.generate(x=x, y=target)


if __name__ == "__main__":
    main()

解决方案 #

1. 迁移ComputeLoss到GPU #

自定义Yolo类中，ComputeLoss实例默认在CPU上初始化，需手动迁移到目标设备：
修改Yolo的__init__方法，添加设备迁移代码，并在实例化时传入设备参数：

def load_model():
    class Yolo(torch.nn.Module):
        def __init__(self, model, device):
            super().__init__()
            self.model = model
            self.model.hyp = {'box': 0.05,
                                'obj': 1.0,
                                'cls': 0.5,
                                'anchor_t': 4.0,
                                'cls_pw': 1.0,
                                'obj_pw': 1.0,
                                'fl_gamma': 0.0
                                }
            self.compute_loss = ComputeLoss(self.model.model.model)
            self.compute_loss.to(device)  # 将损失计算模块迁移到GPU

        def forward(self, x, targets=None):
            if self.training:
                outputs = self.model.model.model(x)
                loss, loss_items = self.compute_loss(outputs, targets)
                loss_components_dict = {"loss_total": loss}
                return loss_components_dict
            else:
                return self.model(x)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = yolov5.load('yolov5s.pt')
    model = Yolo(model, device)  # 传入设备参数
    model.to(device)
    # ... 其余代码不变

2. 同步targets张量设备 #

训练模式下，确保输入x和targets在同一设备上：
修改Yolo的forward方法：

def forward(self, x, targets=None):
    if self.training:
        if targets is not None:
            targets = targets.to(x.device)  # 将targets迁移到x所在设备
        outputs = self.model.model.model(x)
        loss, loss_items = self.compute_loss(outputs, targets)
        loss_components_dict = {"loss_total": loss}
        return loss_components_dict
    else:
        return self.model(x)

3. 统一设备配置参数 #

初始化PyTorchYolo时，直接使用device.type替代字符串判断，避免设备标识不一致：

return PyTorchYolo(model=model,
                    device_type=device.type,
                    input_shape=(3, 640, 640),
                    clip_values=(0, 255), 
                    attack_losses=("loss_total",))

问题根源 #

报错核心是张量设备不匹配：损失计算模块ComputeLoss的参数、训练时的targets张量仍留在CPU，而模型和输入x已迁移到GPU，导致张量拼接操作时触发设备校验错误。

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