基于PyTorch DataParallel的单卡特定操作执行问题咨询

针对你在双GPU上使用DataParallel时遇到的两个特殊操作问题，以下是具体解决方法：

问题1：避免无需批次输入的操作在多卡重复执行 #

DataParallel采用单进程多线程模式，每个GPU对应一个线程。要让操作仅执行一次，只需判断当前线程对应的GPU是否为主设备（默认是cuda:0），仅在主设备上执行该操作即可：

import torch

# 仅在主GPU（cuda:0）执行目标操作
if torch.cuda.current_device() == 0:
    # 这里放只需执行一次的逻辑，比如生成日志、保存中间文件等
    your_single_operation()

这样其他GPU线程会跳过该操作，完全避免冗余耗时。

问题2：仅执行一次即可同步修改所有GPU上的模型 #

DataParallel中，所有GPU上的模型都是model.module（原始模型）的副本。只需在主GPU上修改model.module的参数或属性，后续forward时DataParallel会自动将修改同步到所有GPU的模型副本中：

import torch

if torch.cuda.current_device() == 0:
    # 示例1：修改模型的dropout概率
    model.module.dropout_layer.p = 0.2
    # 示例2：冻结全连接层的参数
    model.module.fc_layer.weight.requires_grad = False
    # 示例3：调整模型的某个超参数
    model.module.custom_param = 100

无需额外的同步操作，下一次模型前向传播时，所有GPU的模型都会使用更新后的状态。

内容的提问来源于stack exchange，提问作者Nagabhushan S N