PyTorch 运行时错误：尽管设备使用一致，但自定义训练循环中的设备不匹配

Question

我正在使用扩散模型架构开发自定义 PyTorch 模型，用于文本到音频生成任务。我的设置涉及使用 BERT 模型处理文本以生成嵌入，然后使用嵌入来调节音频生成。尽管明确地将所有张量和模型移动到同一设备（cuda，如果可用，否则为 cpu），我遇到了一个运行时错误，表明在训练期间设备不匹配。

这是我的设置的简化版本：

import torch
import torch.nn as nn
import librosa
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertModel

# Custom dataset class
class TextAudioDataset(Dataset):
    def __init__(self, audio_dir, text_dir, tokenizer, bert_model, device):
        self.audio_dir = audio_dir
        self.text_dir = text_dir
        self.files = os.listdir(audio_dir)
        self.tokenizer = tokenizer
        self.bert_model = bert_model.to(device)
        self.device = device

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        # Code to load and process audio and text files
        # Returns text_embeddings.to(self.device), audio_tensor.to(self.device)

# Custom model, including a U-Net with FiLM layers conditioned on text embeddings
class DiffusionModel(nn.Module):
    # Model initialization and methods

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
bert_model = BertModel.from_pretrained('bert-base-uncased').to(device)
dataset = TextAudioDataset('./wavs', './transcripts', tokenizer, bert_model, device)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

# Model and training setup
diffusion_model = DiffusionModel(...).to(device)
optimizer = torch.optim.Adam(diffusion_model.parameters(), lr=1e-4)

def train(model, dataloader, optimizer, epochs=10, device=device):
    model.train()
    for epoch in range(epochs):
        for i, (text_embeddings, audio) in enumerate(dataloader):
            text_embeddings = text_embeddings.to(device)
            audio = audio.to(device)
            # Forward pass, loss calculation, backward pass, and optimizer step

train(diffusion_model, dataloader, optimizer, epochs=10, device=device)

尽管将所有张量（text_embeddings、audio）和模型（bert_model、diffusion_model）移动到预期设备，但在训练过程中我仍然收到以下错误：

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

我已经仔细检查了每个张量和模型在涉及它们的任何操作之前是否已明确移动到设备。我的设置或方法中是否有任何我忽略的部分可能导致此设备不匹配问题？

任何有关如何解决此问题的见解或建议将不胜感激！

Answer 1

从这段代码中很难看出。我之前在定制模型时也遇到过类似的问题。当您的自定义模型由多个部分组成时，我发现有时需要将各个部分移动到设备上，因此假设您的模型 A 结合了 LSTM 模型 B 和 CNN 模型 C，而不是您做的

B.to(device)
C.to(device)

而不是

A.to(device)

。如果您属于这种情况，那么值得一试。否则，您需要通过省略某些张量/模型组来进行调试，以找到代码的哪一部分引入了 cpu 张量