El tensor en el modelo es init para cuda al mismo tiempo: notas

Encontré dos problemas:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

Primero veamos la primera pregunta, cómo transferir el tensor de la CPU a la GPU.
El escenario del problema que encontré es el siguiente: defino un modelo con muchas clases y cada clase tiene varios tensores definidos. Antes, simplemente agregaba uno en la CPU donde se informaba el error .cuda(). Era demasiado problemático, así que lo cambié y escribí todas las matrices definidas en funciones o clases:

Esto fue antes:

qtable_8x8 = [[16, 16, 16, 16, 17, 18, 20, 24],
        [16, 16, 16, 17, 18, 20, 24, 25],
        [16, 16, 17, 18, 20, 24, 25, 28],
        [16, 17, 18, 20, 24, 25, 28, 33],
        [17, 18, 20, 24, 25, 28, 33, 41],
        [18, 20, 24, 25, 28, 33, 41, 54],
        [20, 24, 25, 28, 33, 41, 54, 71],
        [24, 25, 28, 33, 41, 54, 71, 91]]

qtable_8x8 = nn.Parameter(torch.from_numpy(np.array(qtable_8x8)), requires_grad=False)

Cámbialo y ponlo en la función:

def def_param():
    qtable_8x8 = [[16, 16, 16, 16, 17, 18, 20, 24],
            [16, 16, 16, 17, 18, 20, 24, 25],
            [16, 16, 17, 18, 20, 24, 25, 28],
            [16, 17, 18, 20, 24, 25, 28, 33],
            [17, 18, 20, 24, 25, 28, 33, 41],
            [18, 20, 24, 25, 28, 33, 41, 54],
            [20, 24, 25, 28, 33, 41, 54, 71],
            [24, 25, 28, 33, 41, 54, 71, 91]]

    qtable_8x8 = nn.Parameter(torch.from_numpy(np.array(qtable_8x8)), requires_grad=False)

La función se coloca en el inicio de la clase:

class quantize(nn.Module):
    def __init__(self, factor):
        super(quantize, self).__init__()
        self.factor = factor
        self.rounding = diff_round
        self.qtable_8x8 = def_param()  # Here
    def forward(self, image):
        image = image.float() * 16 / (self.qtable_8x8 * (self.factor))  # factor
        image = self.rounding(image)
        return image

Al mismo tiempo, si desea que el tensor definido se inicialice en el dispositivo correspondiente junto con la clase, debe cambiar el tensor al formato nn.Parameter().
como sigue:

qtable_8x8 = nn.Parameter(torch.from_numpy(np.array(qtable_8x8)), requires_grad=False)

Tenga en cuenta que si no se requiere gradiente, entoncesrequired_grad=False

class rgb_to_ycbcr_jpeg(nn.Module):
    def __init__(self):
        super(rgb_to_ycbcr_jpeg, self).__init__()
        matrix = np.array(
            [[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5],
             [0.5, -0.418688, -0.081312]], dtype=np.float32).T
        self.shift = nn.Parameter(torch.tensor([0., 128., 128.]), requires_grad=False)
        
        self.matrix = nn.Parameter(torch.from_numpy(matrix), requires_grad=False)

En este punto, to(device)los tensores definidos internamente de .cuda()toda la clase grande DiffIpredse pueden juntar en CUDA.

elif layer == 'H264':
    self.noise_layers = DiffIpred(opt['noise']['H264']['factor']).to(device)

Además, si desea copiar parámetros pasados ​​externamente (como una matriz de parámetros autodefinida) a diferentes tarjetas cuando se ejecutan varias tarjetas en paralelo, debe usar torch.nn.Parameter().

Veamos la segunda pregunta, cómo transferir el tensor de cuda1 a cuda2.
El escenario del problema que encontré es el siguiente: inicialicé varias clases y, al llamar, quería llamar a varias al mismo tiempo, pero no quería escribir demasiados juicios if else, así que usé un diccionario y el El resultado fue la personalización cuando se ejecutaban varias tarjetas en paralelo. Los parámetros de nn.parameter() no se pueden copiar a varias tarjetas.
Así se veía antes del cambio:

class Noise_pool(nn.Module):
    def __init__(self, opt, device):
		super(Noise_pool, self).__init__()
			......... 
			self.Hue           = ColorJitter(opt, distortion='Hue')
			self.Rotation      = kornia.augmentation.RandomRotation(degrees=opt['noise']['Rotation']['degrees'], p=opt['noise']['Rotation']['p'], keepdim=True)
			self.Affine        = kornia.augmentation.RandomAffine(degrees=opt['noise']['Affine']['degrees'], translate=opt['noise']['Affine']['translate'], scale=opt['noise']['Affine']['scale'], shear=opt['noise']['Affine']['shear'], p=opt['noise']['Affine']['p'], keepdim=True)
			#
			self.Cropout       = Cropout(opt)
			self.Dropout       = Dropout(opt)
			# 就是下面这个字典搞的 init的时候并行失败！
			self.noise_pool_in1 = {'Identity': self.Identity, 'JpegTest': self.JpegTest, 'Jpeg': self.Jpeg, \
			                        'JpegMask': self.JpegMask, 'DiffJPEG': self.DiffJPEG, 'Crop': self.Crop, 'Resize': self.Resize, \
			                            'GaussianBlur': self.GaussianBlur, 'Salt_Pepper': self.Salt_Pepper, 'GaussianNoise': self.GaussianNoise, \
			                                'Brightness': self.Brightness, 'Contrast': self.Contrast, 'Saturation': self.Saturation, \
			                                    'Hue': self.Hue, 'Rotation': self.Rotation, 'Affine': self.Affine}
			self.noise_pool_in2 = {'Cropout': self.Cropout, 'Dropout': self.Dropout}
			
	def forward(self, encoded, cover_img, noise_choice):
		......... 

Solución:
Defina 'diccionario' como una función separada y no lo inicialice en init. ... Cuando se use, llámelo directamente hacia adelante. La razón puede ser que la encapsulación dict cambia los atributos de nn.module de la clase, lo que hace que se determine que los parámetros paralelos no se utilizan durante la inicialización.

def forward(self, encoded, cover_img, noise_choice):
    return self.noise_pool_in1()[noise_choice](encoded)
    
def noise_pool_in2(self):
    return {'Cropout': self.Cropout, 'Dropout': self.Dropout}

def noise_pool_in1(self):
	......