如何让我的TensorFlow风格迁移DNN训练脚本在GPU上运行？

嘿，刚好之前做过实时风格迁移的GPU适配，给你整理几个关键步骤，照着改应该就能顺利切换到GPU运行了：

1. 先检测并指定GPU设备 #

首先要确认你的环境能识别到GPU，加这段代码在脚本开头：

import torch

# 自动检测可用GPU，没有的话 fallback 到CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前使用设备: {device}")

2. 把所有模型移到GPU #

不管是你自己实现的FastStyleNet，还是用于损失计算的VGG16，都需要调用.to(device)方法把模型参数迁移到GPU：

import torchvision.models as models

# 实例化风格迁移网络并移到GPU
style_net = FastStyleNet().to(device)

# 加载预训练VGG16（只取特征层，因为我们用它算风格/内容损失）
vgg16 = models.vgg16(pretrained=True).features.to(device)
# 冻结VGG16的参数，因为我们不需要训练它
for param in vgg16.parameters():
    param.requires_grad = False

3. 输入数据必须同步移到GPU #

训练或推理时，每次加载的图像批次（还有风格目标图、内容目标图）都要移到GPU，否则会出现设备不匹配的错误：

# 假设train_loader是你的数据加载器
for batch_idx, (content_imgs, _) in enumerate(train_loader):
    # 把内容图像批次移到GPU
    content_imgs = content_imgs.to(device)
    
    # 风格目标图也要提前预处理后移到GPU（比如只加载一次）
    style_img = preprocess_style_img().to(device)
    
    # 后续前向传播、损失计算全在GPU上进行
    generated_imgs = style_net(content_imgs)
    # ... 计算内容损失、风格损失、总损失

4. 优化器要在模型移GPU后初始化 #

优化器是基于模型参数创建的，所以必须先把模型移到GPU，再初始化优化器，这样优化器才会正确操作GPU上的参数：

# 先移模型到GPU，再初始化优化器
style_net = FastStyleNet().to(device)
optimizer = torch.optim.Adam(style_net.parameters(), lr=1e-3)

常见坑点避坑指南 #

• 如果你的机器有多块GPU，想利用多卡加速，可以用DataParallel包装模型：

  style_net = torch.nn.DataParallel(style_net).to(device)
  

  注意：包装后模型的参数名会带上module.前缀，保存/加载模型时要对应处理（比如加载时用style_net.load_state_dict(torch.load(path), strict=False)）。
• 尽量避免CPU和GPU之间频繁切换张量，这会严重拖慢速度。所有参与计算的张量（包括损失函数里的常量、目标张量）都要确保在同一设备上。
• 如果遇到RuntimeError: Expected all tensors to be on the same device，先排查是不是某个张量（比如损失计算用的权重、临时变量）还留在CPU上，把它也移到GPU就行。

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