Pytorch使用nn.DataParallel失效：只能使用单GPU / 无法使用多GPU训练

官方文档： torch.nn.DataParallel

DataParallel 在 Module 级别实现并行计算，关于使用了 DataParallel 而实际训练时不起作用（只能使用单 GPU）的可能原因：

1. 设置的 Batch Size 小于可用的 GPU 数量，这一点在官方文档中也有说明：Batch Size 需要保证大于使用的 GPU 数量。

2. 未使用 nn.Module 的默认 forward() 方法执行前向传播。

一般情况下，定义的 Model 继承自 torch.nn.Module，在进行前向传播时，调用 model.forward() 函数，能够进行多 GPU 并行训练，一个简单的例子：

# 网络模型定义
class Model(torch.nn.Module):
	def __init__(self):
		super(Model, self).__init__()
		...
	def forward(self, x):
		x = ...
		return x
	
model = Model()   # 初始化网络模型
model = nn.DataParallel(model)   # 指定多GPU
model = model.cuda()
for epoch in range(n_epoch):   #  网络训练过程
	...
	outputs = model(input)
	...

此时，执行 model(input) 调用了默认的 forward() 方法，能够正常进行多 GPU 并行训练。

但是，如果使用自定义 Model 和自定义的方法，则无法进行多 GPU 的并行计算，比如像下面的模型定义：

# 网络模型定义
class Model(torch.nn.Module):
	def __init__(self):
		super(Model, self).__init__()
		...
	def my_forward(self, x): # 自定义forward方法
		x = self.forward(x)
		return x
	def forward(self, x):
		x = ...
		return x

model = Model()   # 初始化网络模型
model = nn.DataParallel(model)   # 指定多GPU
model = model.cuda()
for epoch in range(n_epoch):   #  网络训练过程
	...
	outputs = model.my_forward(input)  # 使用自定义的forward方法
	...

在使用了 model = nn.DataParallel(model) 之后，可以发现 model 变成了 DataParallel 类型的变量，其中 module 变量才是真正的网络模型，可以通过 model.module 获取。

此时，如果调用自定义的方法 model.my_forward(input)，会报错：

'DataParallel' object has no attribute 'my_forward'

要解决这个 bug 很简单，只需要利用 model = model.module 把 model 从 DataParallel 里拎出来即可：

if isinstance(model, torch.nn.DataParallel):
    model = model.module

这样做虽然能够正常训练，但同时也意味着无法使用 DataParallel 进行多 GPU 并行训练，也就出现了使用了 DataParallel，但无效的问题。

那么，如果 my_forward() 方法非定义不可呢？尝试曲线救国，将自定义方法单独拎出来自成一类，然后在该类中调用默认的 forward() 方法。举个例子：

class Model(torch.nn.Module):
	def __init__(self):
		super(Model, self).__init__()
		...
	def forward(self, x):
		x = ...
		return x

# 新定义类
class MyForward(torch.nn.Module):
    def __init__(self):
        super().__init__()
        ...
    def my_forward(self, x):
    	...
        x = model(x)
        return x

model = Model()   # 初始化网络模型
my_forward = MyForward()   # 初始化新定义类
model = nn.DataParallel(model)
model = model.cuda()
for epoch in range(n_epoch):
	...
	outputs = my_forward(input)
	...

至此，问题解决 ~