如何基于Dataset创建保留requires_grad=True属性的DataLoader

问题原因 #

你的猜测是正确的，PyTorch的DataLoader默认在拼接batch、或者多进程数据传输的过程中会生成新的张量副本，默认不会继承原张量的requires_grad属性，所以会出现属性被重置为False的情况。

最优解决方案 #

推荐按优先级选择以下方案：

方案1：拿到batch后手动开启梯度（最推荐） #

这是PyTorch官方推荐的处理方式，数据加载和梯度计算逻辑分离，兼容所有加载配置（包括多进程加载），实现成本最低：

it = iter(trainDataLoader)
first = next(it)
# 手动为batch张量开启梯度追踪
batch_data = first[0].requires_grad_()
batch_label = first[1].requires_grad_()

# 验证梯度属性
print(batch_data[0].requires_grad) # 输出True

方案2：自定义collate_fn保留梯度属性 #

如果你需要DataLoader直接输出带梯度的张量，可以替换默认的拼接函数，同时要注意必须设置num_workers=0（多进程序列化传输会强制丢失梯度属性）：

def custom_collate(batch):
    # 手动堆叠样本，保留梯度属性
    data = torch.stack([sample[0] for sample in batch])
    labels = torch.stack([sample[1] for sample in batch])
    return data, labels

# 初始化DataLoader时指定自定义collate_fn
trainDataLoader = DataLoader(
    dataset=_data, 
    batch_size=15, 
    shuffle=True,
    collate_fn=custom_collate,
    num_workers=0 # 必须设为0，多进程场景无法保留梯度
)

注意事项 #

除非特殊场景需要在数据加载阶段就保留梯度，否则优先选择方案1，不会引入额外的兼容问题，性能也不会有损失。

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