PyTorch自定义训练循环中设备不匹配问题排查求助

我基于扩散模型架构开发自定义PyTorch文本转音频生成模型，采用BERT模型处理文本生成嵌入向量以指导音频生成。已将所有张量、模型显式迁移至同一设备（优先CUDA，否则CPU），但训练时仍触发设备不匹配的RuntimeError。

简化代码实现 #

import torch
import torch.nn as nn
import librosa
import os  # 原代码遗漏该导入，需补充
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertModel

# 自定义数据集类
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):
        # 加载并处理音频和文本文件的代码
        # 返回 text_embeddings.to(self.device), audio_tensor.to(self.device)

# 自定义模型，包含基于文本嵌入条件的FiLM层U-Net
class DiffusionModel(nn.Module):
    # 模型初始化和方法

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)

# 模型和训练设置
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)
            # 前向传播、损失计算、反向传播、优化器更新步骤

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

错误信息 #

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

排查关键点 #

• DataLoader默认行为导致张量回传CPU：即使在__getitem__中将张量移到GPU，DataLoader的默认collate_fn会把分散的GPU张量拉回CPU拼接成batch。解决方法：在__getitem__中返回CPU张量，训练循环内统一移到device；或自定义collate_fn保持张量在GPU上（大batch时慎用，易显存不足）。
• BERT推理的输入张量未移设备：__getitem__中用BERT生成嵌入时，需确保tokenizer输出的input_ids、attention_mask等张量先移到目标设备，否则BERT会在CPU上生成嵌入，导致后续设备不匹配。
• DiffusionModel内部固定张量未指定设备：检查模型初始化代码，若存在手动创建的固定张量（如时间步嵌入、FiLM层参数），需显式指定device=device，否则默认在CPU上。
• 优化器创建时机二次确认：确保在模型移到device之后再创建优化器，否则优化器会绑定CPU上的模型参数（你的代码中该步骤正确，可作为验证项）。

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