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)