Support factory kwargs in torch.nn modules (pytorch#54508)

Summary:
Continuation of pytorch#53144

Pull Request resolved: pytorch#54508

Reviewed By: albanD

Differential Revision: D27939544

Pulled By: jbschlosser

fbshipit-source-id: 4bf517e5f74f093e27ca38a85e732da65e44d805

Author: jbschlosser

test/run_test.py

@@ -65,6 +65,7 @@
     'test_linalg',
     'test_logging',
     'test_mkldnn',
+    'test_module_init',
     'test_multiprocessing',
     'test_multiprocessing_spawn',
     'distributed/test_nccl',

test/test_module_init.py

@@ -3,3 +3,44 @@
 from .parallel import DataParallel
 from . import init
 from . import utils
+
+
+def factory_kwargs(kwargs):
+    r"""
+    Given kwargs, returns a canonicalized dict of factory kwargs that can be directly passed
+    to factory functions like torch.empty, or errors if unrecognized kwargs are present.
+
+    This function makes it simple to write code like this::
+
+        class MyModule(nn.Module):
+            def __init__(self, **kwargs):
+                factory_kwargs = torch.nn.factory_kwargs(kwargs)
+                self.weight = Parameter(torch.empty(10, **factory_kwargs))
+
+    Why should you use this function instead of just passing `kwargs` along directly?
+
+    1. This function does error validation, so if there are unexpected kwargs we will
+    immediately report an error, instead of deferring it to the factory call
+    2. This function supports a special `factory_kwargs` argument, which can be used to
+    explicitly specify a kwarg to be used for factory functions, in the event one of the
+    factory kwargs conflicts with an already existing argument in the signature (e.g.
+    in the signature ``def f(dtype, **kwargs)``, you can specify ``dtype`` for factory
+    functions, as distinct from the dtype argument, by saying
+    ``f(dtype1, factory_kwargs={"dtype": dtype2})``)
+    """
+    if kwargs is None:
+        return {}
+    simple_keys = {"device", "dtype", "memory_format"}
+    expected_keys = simple_keys | {"factory_kwargs"}
+    if not kwargs.keys() <= expected_keys:
+        raise TypeError(f"unexpected kwargs {kwargs.keys() - expected_keys}")
+
+    # guarantee no input kwargs is untouched
+    r = dict(kwargs.get("factory_kwargs", {}))
+    for k in simple_keys:
+        if k in kwargs:
+            if k in r:
+                raise TypeError(f"{k} specified twice, in **kwargs and in factory_kwargs")
+            r[k] = kwargs[k]
+
+    return r

torch/nn/__init__.py

@@ -872,7 +872,8 @@ class MultiheadAttention(Module):
     bias_v: Optional[torch.Tensor]
 
     def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False,
-                 kdim=None, vdim=None, batch_first=False):
+                 kdim=None, vdim=None, batch_first=False, device=None, dtype=None) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
         super(MultiheadAttention, self).__init__()
         self.embed_dim = embed_dim
         self.kdim = kdim if kdim is not None else embed_dim
@@ -886,25 +887,25 @@ def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=Fals
         assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
 
         if self._qkv_same_embed_dim is False:
-            self.q_proj_weight = Parameter(torch.empty(embed_dim, embed_dim))
-            self.k_proj_weight = Parameter(torch.empty(embed_dim, self.kdim))
-            self.v_proj_weight = Parameter(torch.empty(embed_dim, self.vdim))
+            self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim), **factory_kwargs))
+            self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim), **factory_kwargs))
+            self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim), **factory_kwargs))
             self.register_parameter('in_proj_weight', None)
         else:
-            self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
+            self.in_proj_weight = Parameter(torch.empty((3 * embed_dim, embed_dim), **factory_kwargs))
             self.register_parameter('q_proj_weight', None)
             self.register_parameter('k_proj_weight', None)
             self.register_parameter('v_proj_weight', None)
 
         if bias:
-            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
+            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim, **factory_kwargs))
         else:
             self.register_parameter('in_proj_bias', None)
-        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)
 
         if add_bias_kv:
-            self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
-            self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
+            self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
+            self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
         else:
             self.bias_k = self.bias_v = None
 
@@ -1057,10 +1058,12 @@ class PReLU(Module):
     __constants__ = ['num_parameters']
     num_parameters: int
 
-    def __init__(self, num_parameters: int = 1, init: float = 0.25) -> None:
+    def __init__(self, num_parameters: int = 1, init: float = 0.25,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
         self.num_parameters = num_parameters
         super(PReLU, self).__init__()
-        self.weight = Parameter(torch.empty(num_parameters).fill_(init))
+        self.weight = Parameter(torch.empty(num_parameters, **factory_kwargs).fill_(init))
 
     def forward(self, input: Tensor) -> Tensor:
         return F.prelu(input, self.weight)

torch/nn/modules/activation.py

@@ -115,8 +115,11 @@ def __init__(
         n_classes: int,
         cutoffs: Sequence[int],
         div_value: float = 4.,
-        head_bias: bool = False
+        head_bias: bool = False,
+        device=None,
+        dtype=None
     ) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
         super(AdaptiveLogSoftmaxWithLoss, self).__init__()
 
         cutoffs = list(cutoffs)
@@ -141,7 +144,8 @@ def __init__(
         self.n_clusters = len(self.cutoffs) - 1
         self.head_size = self.shortlist_size + self.n_clusters
 
-        self.head = Linear(self.in_features, self.head_size, bias=self.head_bias)
+        self.head = Linear(self.in_features, self.head_size, bias=self.head_bias,
+                           **factory_kwargs)
         self.tail = ModuleList()
 
         for i in range(self.n_clusters):
@@ -150,8 +154,8 @@ def __init__(
             osz = self.cutoffs[i + 1] - self.cutoffs[i]
 
             projection = Sequential(
-                Linear(self.in_features, hsz, bias=False),
-                Linear(hsz, osz, bias=False)
+                Linear(self.in_features, hsz, bias=False, **factory_kwargs),
+                Linear(hsz, osz, bias=False, **factory_kwargs),
             )
 
             self.tail.append(projection)

torch/nn/modules/adaptive.py

@@ -31,25 +31,28 @@ def __init__(
         momentum: float = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device=None,
+        dtype=None
     ) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
         super(_NormBase, self).__init__()
         self.num_features = num_features
         self.eps = eps
         self.momentum = momentum
         self.affine = affine
         self.track_running_stats = track_running_stats
         if self.affine:
-            self.weight = Parameter(torch.empty(num_features))
-            self.bias = Parameter(torch.empty(num_features))
+            self.weight = Parameter(torch.empty(num_features, **factory_kwargs))
+            self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
         else:
             self.register_parameter("weight", None)
             self.register_parameter("bias", None)
         if self.track_running_stats:
-            self.register_buffer("running_mean", torch.zeros(num_features))
-            self.register_buffer("running_var", torch.ones(num_features))
-            self.register_buffer(
-                "num_batches_tracked", torch.tensor(0, dtype=torch.long)
-            )
+            self.register_buffer('running_mean', torch.zeros(num_features, **factory_kwargs))
+            self.register_buffer('running_var', torch.ones(num_features, **factory_kwargs))
+            self.register_buffer('num_batches_tracked',
+                                 torch.tensor(0, dtype=torch.long,
+                                              **{k: v for k, v in factory_kwargs.items() if k != 'dtype'}))
         else:
             self.register_buffer("running_mean", None)
             self.register_buffer("running_var", None)
@@ -117,9 +120,12 @@ def __init__(
         momentum=0.1,
         affine=True,
         track_running_stats=True,
+        device=None,
+        dtype=None
     ):
+        factory_kwargs = {'device': device, 'dtype': dtype}
         super(_BatchNorm, self).__init__(
-            num_features, eps, momentum, affine, track_running_stats
+            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
         )
 
     def forward(self, input: Tensor) -> Tensor:
@@ -178,7 +184,9 @@ class _LazyBatchNorm(LazyModuleMixin, _BatchNorm):
     weight: UninitializedParameter  # type: ignore[assignment]
     bias: UninitializedParameter  # type: ignore[assignment]
 
-    def __init__(self, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
+    def __init__(self, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True,
+                 device=None, dtype=None) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
         super(_LazyBatchNorm, self).__init__(
             # affine and track_running_stats are hardcoded to False to
             # avoid creating tensors that will soon be overwritten.
@@ -187,16 +195,18 @@ def __init__(self, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True
             momentum,
             False,
             False,
+            **factory_kwargs,
         )
         self.affine = affine
         self.track_running_stats = track_running_stats
         if self.affine:
-            self.weight = UninitializedParameter()
-            self.bias = UninitializedParameter()
+            self.weight = UninitializedParameter(**factory_kwargs)
+            self.bias = UninitializedParameter(**factory_kwargs)
         if self.track_running_stats:
-            self.running_mean = UninitializedBuffer()
-            self.running_var = UninitializedBuffer()
-            self.num_batches_tracked = torch.tensor(0, dtype=torch.long)
+            self.running_mean = UninitializedBuffer(**factory_kwargs)
+            self.running_var = UninitializedBuffer(**factory_kwargs)
+            self.num_batches_tracked = torch.tensor(
+                0, dtype=torch.long, **{k: v for k, v in factory_kwargs.items() if k != 'dtype'})
 
     def reset_parameters(self) -> None:
         if not self.has_uninitialized_params() and self.num_features != 0:
@@ -640,9 +650,12 @@ def __init__(
         affine: bool = True,
         track_running_stats: bool = True,
         process_group: Optional[Any] = None,
+        device=None,
+        dtype=None
     ) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
         super(SyncBatchNorm, self).__init__(
-            num_features, eps, momentum, affine, track_running_stats
+            num_features, eps, momentum, affine, track_running_stats, **factory_kwargs
         )
         self.process_group = process_group