为啥 PyTorch nn.Module.cuda() 不移动模块张量而只移动参数和缓冲区到 GPU?
Posted
技术标签:
【中文标题】为啥 PyTorch nn.Module.cuda() 不移动模块张量而只移动参数和缓冲区到 GPU?【英文标题】:Why PyTorch nn.Module.cuda() not moving Module tensor but only parameters and buffers to GPU?为什么 PyTorch nn.Module.cuda() 不移动模块张量而只移动参数和缓冲区到 GPU? 【发布时间】:2020-07-09 12:57:06 【问题描述】:nn.Module.cuda()
将所有模型参数和缓冲区移动到 GPU。
但为什么不是模型成员张量?
class ToyModule(torch.nn.Module):
def __init__(self) -> None:
super(ToyModule, self).__init__()
self.layer = torch.nn.Linear(2, 2)
self.expected_moved_cuda_tensor = torch.tensor([0, 2, 3])
def forward(self, input: torch.Tensor) -> torch.Tensor:
return self.layer(input)
toy_module = ToyModule()
toy_module.cuda()
next(toy_module.layer.parameters()).device
>>> device(type='cuda', index=0)
对于模型成员张量,设备保持不变。
>>> toy_module.expected_moved_cuda_tensor.device
device(type='cpu')
【问题讨论】:
【参考方案1】:如果您在模块内定义张量,则需要将其注册为参数或缓冲区,以便模块知道它。
参数 是要训练的张量,将由model.parameters()
返回。它们很容易注册,您只需将张量包装在 nn.Parameter
类型中,它将自动注册。请注意,只有浮点张量可以作为参数。
class ToyModule(torch.nn.Module):
def __init__(self) -> None:
super(ToyModule, self).__init__()
self.layer = torch.nn.Linear(2, 2)
# registering expected_moved_cuda_tensor as a trainable parameter
self.expected_moved_cuda_tensor = torch.nn.Parameter(torch.tensor([0., 2., 3.]))
def forward(self, input: torch.Tensor) -> torch.Tensor:
return self.layer(input)
缓冲区 是将在模块中注册的张量,因此.cuda()
之类的方法会影响它们,但model.parameters()
不会返回它们。缓冲区不限于特定的数据类型。
class ToyModule(torch.nn.Module):
def __init__(self) -> None:
super(ToyModule, self).__init__()
self.layer = torch.nn.Linear(2, 2)
# registering expected_moved_cuda_tensor as a buffer
# Note: this creates a new member variable named expected_moved_cuda_tensor
self.register_buffer('expected_moved_cuda_tensor', torch.tensor([0, 2, 3])))
def forward(self, input: torch.Tensor) -> torch.Tensor:
return self.layer(input)
在上述两种情况下,以下代码的行为相同
>>> toy_module = ToyModule()
>>> toy_module.cuda()
>>> next(toy_module.layer.parameters()).device
device(type='cuda', index=0)
>>> toy_module.expected_moved_cuda_tensor.device
device(type='cuda', index=0)
【讨论】:
以上是关于为啥 PyTorch nn.Module.cuda() 不移动模块张量而只移动参数和缓冲区到 GPU?的主要内容,如果未能解决你的问题,请参考以下文章
为啥这个 Pytorch 官方教程中没有 .train() 方法?