如果您在模块内定义张量,则需要将其注册为参数或缓冲区,以便模块知道它。
参数是要训练的张量,并将通过以下方式返回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)
Buffers是将在模块中注册的张量,因此方法如下.cuda()
会影响他们,但他们会not被返回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)