Python PyTorch skip_init用法及代码示例

本文简要介绍python语言中 torch.nn.utils.skip_init 的用法。

用法:

torch.nn.utils.skip_init(module_cls, *args, **kwargs)

参数：

• module_cls-类对象；应该是 torch.nn.Module 的子类

• args-args 传递给模块的构造函数

• kwargs-kwargs 传递给模块的构造函数

返回：

具有未初始化参数/缓冲区的实例化模块

给定一个模块类对象和 args /kwargs，在不初始化参数 /缓冲区的情况下实例化模块。如果初始化很慢或将执行自定义初始化，这可能很有用，从而不需要默认初始化。由于此函数的实现方式，对此有一些警告：

1. 模块必须在其构造函数中接受 device 参数，该参数传递给构造期间创建的任何参数或缓冲区。

2. 除了初始化(即来自torch.nn.init的函数)之外，模块不得对其构造函数中的参数执行任何计算。

如果满足这些条件，则可以使用未初始化的参数/缓冲区值来实例化模块，就像使用 torch.empty() 创建一样。

例子：

>>> import torch
>>> m = torch.nn.utils.skip_init(torch.nn.Linear, 5, 1)
>>> m.weight
Parameter containing:
tensor([[0.0000e+00, 1.5846e+29, 7.8307e+00, 2.5250e-29, 1.1210e-44]],
       requires_grad=True)
>>> m2 = torch.nn.utils.skip_init(torch.nn.Linear, in_features=6, out_features=1)
>>> m2.weight
Parameter containing:
tensor([[-1.4677e+24,  4.5915e-41,  1.4013e-45,  0.0000e+00, -1.4677e+24,
          4.5915e-41]], requires_grad=True)