| from enum import Enum, auto |
| |
| import torch |
| from torch import Tensor |
| from ..utils import parametrize |
| from ..modules import Module |
| from .. import functional as F |
| |
| from typing import Optional |
| |
| __all__ = ['orthogonal', 'spectral_norm', 'weight_norm'] |
| |
| |
| def _is_orthogonal(Q, eps=None): |
| n, k = Q.size(-2), Q.size(-1) |
| Id = torch.eye(k, dtype=Q.dtype, device=Q.device) |
| # A reasonable eps, but not too large |
| eps = 10. * n * torch.finfo(Q.dtype).eps |
| return torch.allclose(Q.mH @ Q, Id, atol=eps) |
| |
| |
| def _make_orthogonal(A): |
| """ Assume that A is a tall matrix. |
| Compute the Q factor s.t. A = QR (A may be complex) and diag(R) is real and non-negative |
| """ |
| X, tau = torch.geqrf(A) |
| Q = torch.linalg.householder_product(X, tau) |
| # The diagonal of X is the diagonal of R (which is always real) so we normalise by its signs |
| Q *= X.diagonal(dim1=-2, dim2=-1).sgn().unsqueeze(-2) |
| return Q |
| |
| |
| class _OrthMaps(Enum): |
| matrix_exp = auto() |
| cayley = auto() |
| householder = auto() |
| |
| |
| class _Orthogonal(Module): |
| base: Tensor |
| |
| def __init__(self, |
| weight, |
| orthogonal_map: _OrthMaps, |
| *, |
| use_trivialization=True) -> None: |
| super().__init__() |
| |
| # Note [Householder complex] |
| # For complex tensors, it is not possible to compute the tensor `tau` necessary for |
| # linalg.householder_product from the reflectors. |
| # To see this, note that the reflectors have a shape like: |
| # 0 0 0 |
| # * 0 0 |
| # * * 0 |
| # which, for complex matrices, give n(n-1) (real) parameters. Now, you need n^2 parameters |
| # to parametrize the unitary matrices. Saving tau on its own does not work either, because |
| # not every combination of `(A, tau)` gives a unitary matrix, meaning that if we optimise |
| # them as independent tensors we would not maintain the constraint |
| # An equivalent reasoning holds for rectangular matrices |
| if weight.is_complex() and orthogonal_map == _OrthMaps.householder: |
| raise ValueError("The householder parametrization does not support complex tensors.") |
| |
| self.shape = weight.shape |
| self.orthogonal_map = orthogonal_map |
| if use_trivialization: |
| self.register_buffer("base", None) |
| |
| def forward(self, X: torch.Tensor) -> torch.Tensor: |
| n, k = X.size(-2), X.size(-1) |
| transposed = n < k |
| if transposed: |
| X = X.mT |
| n, k = k, n |
| # Here n > k and X is a tall matrix |
| if self.orthogonal_map == _OrthMaps.matrix_exp or self.orthogonal_map == _OrthMaps.cayley: |
| # We just need n x k - k(k-1)/2 parameters |
| X = X.tril() |
| if n != k: |
| # Embed into a square matrix |
| X = torch.cat([X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1) |
| A = X - X.mH |
| # A is skew-symmetric (or skew-hermitian) |
| if self.orthogonal_map == _OrthMaps.matrix_exp: |
| Q = torch.matrix_exp(A) |
| elif self.orthogonal_map == _OrthMaps.cayley: |
| # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1} |
| Id = torch.eye(n, dtype=A.dtype, device=A.device) |
| Q = torch.linalg.solve(torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5)) |
| # Q is now orthogonal (or unitary) of size (..., n, n) |
| if n != k: |
| Q = Q[..., :k] |
| # Q is now the size of the X (albeit perhaps transposed) |
| else: |
| # X is real here, as we do not support householder with complex numbers |
| A = X.tril(diagonal=-1) |
| tau = 2. / (1. + (A * A).sum(dim=-2)) |
| Q = torch.linalg.householder_product(A, tau) |
| # The diagonal of X is 1's and -1's |
| # We do not want to differentiate through this or update the diagonal of X hence the casting |
| Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2) |
| |
| if hasattr(self, "base"): |
| Q = self.base @ Q |
| if transposed: |
| Q = Q.mT |
| return Q |
| |
| @torch.autograd.no_grad() |
| def right_inverse(self, Q: torch.Tensor) -> torch.Tensor: |
| if Q.shape != self.shape: |
| raise ValueError(f"Expected a matrix or batch of matrices of shape {self.shape}. " |
| f"Got a tensor of shape {Q.shape}.") |
| |
| Q_init = Q |
| n, k = Q.size(-2), Q.size(-1) |
| transpose = n < k |
| if transpose: |
| Q = Q.mT |
| n, k = k, n |
| |
| # We always make sure to always copy Q in every path |
| if not hasattr(self, "base"): |
| # Note [right_inverse expm cayley] |
| # If we do not have use_trivialization=True, we just implement the inverse of the forward |
| # map for the Householder. To see why, think that for the Cayley map, |
| # we would need to find the matrix X \in R^{n x k} such that: |
| # Y = torch.cat([X.tril(), X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1) |
| # A = Y - Y.mH |
| # cayley(A)[:, :k] |
| # gives the original tensor. It is not clear how to do this. |
| # Perhaps via some algebraic manipulation involving the QR like that of |
| # Corollary 2.2 in Edelman, Arias and Smith? |
| if self.orthogonal_map == _OrthMaps.cayley or self.orthogonal_map == _OrthMaps.matrix_exp: |
| raise NotImplementedError("It is not possible to assign to the matrix exponential " |
| "or the Cayley parametrizations when use_trivialization=False.") |
| |
| # If parametrization == _OrthMaps.householder, make Q orthogonal via the QR decomposition. |
| # Here Q is always real because we do not support householder and complex matrices. |
| # See note [Householder complex] |
| A, tau = torch.geqrf(Q) |
| # We want to have a decomposition X = QR with diag(R) > 0, as otherwise we could |
| # decompose an orthogonal matrix Q as Q = (-Q)@(-Id), which is a valid QR decomposition |
| # The diagonal of Q is the diagonal of R from the qr decomposition |
| A.diagonal(dim1=-2, dim2=-1).sign_() |
| # Equality with zero is ok because LAPACK returns exactly zero when it does not want |
| # to use a particular reflection |
| A.diagonal(dim1=-2, dim2=-1)[tau == 0.] *= -1 |
| return A.mT if transpose else A |
| else: |
| if n == k: |
| # We check whether Q is orthogonal |
| if not _is_orthogonal(Q): |
| Q = _make_orthogonal(Q) |
| else: # Is orthogonal |
| Q = Q.clone() |
| else: |
| # Complete Q into a full n x n orthogonal matrix |
| N = torch.randn(*(Q.size()[:-2] + (n, n - k)), dtype=Q.dtype, device=Q.device) |
| Q = torch.cat([Q, N], dim=-1) |
| Q = _make_orthogonal(Q) |
| self.base = Q |
| |
| # It is necessary to return the -Id, as we use the diagonal for the |
| # Householder parametrization. Using -Id makes: |
| # householder(torch.zeros(m,n)) == torch.eye(m,n) |
| # Poor man's version of eye_like |
| neg_Id = torch.zeros_like(Q_init) |
| neg_Id.diagonal(dim1=-2, dim2=-1).fill_(-1.) |
| return neg_Id |
| |
| |
| def orthogonal(module: Module, |
| name: str = 'weight', |
| orthogonal_map: Optional[str] = None, |
| *, |
| use_trivialization: bool = True) -> Module: |
| r"""Applies an orthogonal or unitary parametrization to a matrix or a batch of matrices. |
| |
| Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, the parametrized |
| matrix :math:`Q \in \mathbb{K}^{m \times n}` is **orthogonal** as |
| |
| .. math:: |
| |
| \begin{align*} |
| Q^{\text{H}}Q &= \mathrm{I}_n \mathrlap{\qquad \text{if }m \geq n}\\ |
| QQ^{\text{H}} &= \mathrm{I}_m \mathrlap{\qquad \text{if }m < n} |
| \end{align*} |
| |
| where :math:`Q^{\text{H}}` is the conjugate transpose when :math:`Q` is complex |
| and the transpose when :math:`Q` is real-valued, and |
| :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix. |
| In plain words, :math:`Q` will have orthonormal columns whenever :math:`m \geq n` |
| and orthonormal rows otherwise. |
| |
| If the tensor has more than two dimensions, we consider it as a batch of matrices of shape `(..., m, n)`. |
| |
| The matrix :math:`Q` may be parametrized via three different ``orthogonal_map`` in terms of the original tensor: |
| |
| - ``"matrix_exp"``/``"cayley"``: |
| the :func:`~torch.matrix_exp` :math:`Q = \exp(A)` and the `Cayley map`_ |
| :math:`Q = (\mathrm{I}_n + A/2)(\mathrm{I}_n - A/2)^{-1}` are applied to a skew-symmetric |
| :math:`A` to give an orthogonal matrix. |
| - ``"householder"``: computes a product of Householder reflectors |
| (:func:`~torch.linalg.householder_product`). |
| |
| ``"matrix_exp"``/``"cayley"`` often make the parametrized weight converge faster than |
| ``"householder"``, but they are slower to compute for very thin or very wide matrices. |
| |
| If ``use_trivialization=True`` (default), the parametrization implements the "Dynamic Trivialization Framework", |
| where an extra matrix :math:`B \in \mathbb{K}^{n \times n}` is stored under |
| ``module.parametrizations.weight[0].base``. This helps the |
| convergence of the parametrized layer at the expense of some extra memory use. |
| See `Trivializations for Gradient-Based Optimization on Manifolds`_ . |
| |
| Initial value of :math:`Q`: |
| If the original tensor is not parametrized and ``use_trivialization=True`` (default), the initial value |
| of :math:`Q` is that of the original tensor if it is orthogonal (or unitary in the complex case) |
| and it is orthogonalized via the QR decomposition otherwise (see :func:`torch.linalg.qr`). |
| Same happens when it is not parametrized and ``orthogonal_map="householder"`` even when ``use_trivialization=False``. |
| Otherwise, the initial value is the result of the composition of all the registered |
| parametrizations applied to the original tensor. |
| |
| .. note:: |
| This function is implemented using the parametrization functionality |
| in :func:`~torch.nn.utils.parametrize.register_parametrization`. |
| |
| |
| .. _`Cayley map`: https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map |
| .. _`Trivializations for Gradient-Based Optimization on Manifolds`: https://arxiv.org/abs/1909.09501 |
| |
| Args: |
| module (nn.Module): module on which to register the parametrization. |
| name (str, optional): name of the tensor to make orthogonal. Default: ``"weight"``. |
| orthogonal_map (str, optional): One of the following: ``"matrix_exp"``, ``"cayley"``, ``"householder"``. |
| Default: ``"matrix_exp"`` if the matrix is square or complex, ``"householder"`` otherwise. |
| use_trivialization (bool, optional): whether to use the dynamic trivialization framework. |
| Default: ``True``. |
| |
| Returns: |
| The original module with an orthogonal parametrization registered to the specified |
| weight |
| |
| Example:: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) |
| >>> orth_linear = orthogonal(nn.Linear(20, 40)) |
| >>> orth_linear |
| ParametrizedLinear( |
| in_features=20, out_features=40, bias=True |
| (parametrizations): ModuleDict( |
| (weight): ParametrizationList( |
| (0): _Orthogonal() |
| ) |
| ) |
| ) |
| >>> # xdoctest: +IGNORE_WANT |
| >>> Q = orth_linear.weight |
| >>> torch.dist(Q.T @ Q, torch.eye(20)) |
| tensor(4.9332e-07) |
| """ |
| weight = getattr(module, name, None) |
| if not isinstance(weight, Tensor): |
| raise ValueError( |
| f"Module '{module}' has no parameter or buffer with name '{name}'" |
| ) |
| |
| # We could implement this for 1-dim tensors as the maps on the sphere |
| # but I believe it'd bite more people than it'd help |
| if weight.ndim < 2: |
| raise ValueError("Expected a matrix or batch of matrices. " |
| f"Got a tensor of {weight.ndim} dimensions.") |
| |
| if orthogonal_map is None: |
| orthogonal_map = "matrix_exp" if weight.size(-2) == weight.size(-1) or weight.is_complex() else "householder" |
| |
| orth_enum = getattr(_OrthMaps, orthogonal_map, None) |
| if orth_enum is None: |
| raise ValueError('orthogonal_map has to be one of "matrix_exp", "cayley", "householder". ' |
| f'Got: {orthogonal_map}') |
| orth = _Orthogonal(weight, |
| orth_enum, |
| use_trivialization=use_trivialization) |
| parametrize.register_parametrization(module, name, orth, unsafe=True) |
| return module |
| |
| |
| class _WeightNorm(Module): |
| def __init__( |
| self, |
| dim: Optional[int] = 0, |
| ) -> None: |
| super().__init__() |
| if dim is None: |
| dim = -1 |
| self.dim = dim |
| |
| def forward(self, weight_g, weight_v): |
| return torch._weight_norm(weight_v, weight_g, self.dim) |
| |
| def right_inverse(self, weight): |
| weight_g = torch.norm_except_dim(weight, 2, self.dim) |
| weight_v = weight |
| |
| return weight_g, weight_v |
| |
| |
| def weight_norm(module: Module, name: str = 'weight', dim: int = 0): |
| r"""Applies weight normalization to a parameter in the given module. |
| |
| .. math:: |
| \mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|} |
| |
| Weight normalization is a reparameterization that decouples the magnitude |
| of a weight tensor from its direction. This replaces the parameter specified |
| by :attr:`name` with two parameters: one specifying the magnitude |
| and one specifying the direction. |
| |
| By default, with ``dim=0``, the norm is computed independently per output |
| channel/plane. To compute a norm over the entire weight tensor, use |
| ``dim=None``. |
| |
| See https://arxiv.org/abs/1602.07868 |
| |
| Args: |
| module (Module): containing module |
| name (str, optional): name of weight parameter |
| dim (int, optional): dimension over which to compute the norm |
| |
| Returns: |
| The original module with the weight norm hook |
| |
| Example:: |
| |
| >>> m = weight_norm(nn.Linear(20, 40), name='weight') |
| >>> m |
| ParametrizedLinear( |
| in_features=20, out_features=40, bias=True |
| (parametrizations): ModuleDict( |
| (weight): ParametrizationList( |
| (0): _WeightNorm() |
| ) |
| ) |
| ) |
| >>> m.parametrizations.weight.original0.size() |
| torch.Size([40, 1]) |
| >>> m.parametrizations.weight.original1.size() |
| torch.Size([40, 20]) |
| |
| """ |
| _weight_norm = _WeightNorm(dim) |
| parametrize.register_parametrization(module, name, _weight_norm, unsafe=True) |
| |
| def _weight_norm_compat_hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): |
| g_key = f"{prefix}{name}_g" |
| v_key = f"{prefix}{name}_v" |
| if g_key in state_dict and v_key in state_dict: |
| original0 = state_dict.pop(g_key) |
| original1 = state_dict.pop(v_key) |
| state_dict[f"{prefix}parametrizations.{name}.original0"] = original0 |
| state_dict[f"{prefix}parametrizations.{name}.original1"] = original1 |
| module._register_load_state_dict_pre_hook(_weight_norm_compat_hook) |
| return module |
| |
| |
| class _SpectralNorm(Module): |
| def __init__( |
| self, |
| weight: torch.Tensor, |
| n_power_iterations: int = 1, |
| dim: int = 0, |
| eps: float = 1e-12 |
| ) -> None: |
| super().__init__() |
| ndim = weight.ndim |
| if dim >= ndim or dim < -ndim: |
| raise IndexError("Dimension out of range (expected to be in range of " |
| f"[-{ndim}, {ndim - 1}] but got {dim})") |
| |
| if n_power_iterations <= 0: |
| raise ValueError('Expected n_power_iterations to be positive, but ' |
| f'got n_power_iterations={n_power_iterations}') |
| self.dim = dim if dim >= 0 else dim + ndim |
| self.eps = eps |
| if ndim > 1: |
| # For ndim == 1 we do not need to approximate anything (see _SpectralNorm.forward) |
| self.n_power_iterations = n_power_iterations |
| weight_mat = self._reshape_weight_to_matrix(weight) |
| h, w = weight_mat.size() |
| |
| u = weight_mat.new_empty(h).normal_(0, 1) |
| v = weight_mat.new_empty(w).normal_(0, 1) |
| self.register_buffer('_u', F.normalize(u, dim=0, eps=self.eps)) |
| self.register_buffer('_v', F.normalize(v, dim=0, eps=self.eps)) |
| |
| # Start with u, v initialized to some reasonable values by performing a number |
| # of iterations of the power method |
| self._power_method(weight_mat, 15) |
| |
| def _reshape_weight_to_matrix(self, weight: torch.Tensor) -> torch.Tensor: |
| # Precondition |
| assert weight.ndim > 1 |
| |
| if self.dim != 0: |
| # permute dim to front |
| weight = weight.permute(self.dim, *(d for d in range(weight.dim()) if d != self.dim)) |
| |
| return weight.flatten(1) |
| |
| @torch.autograd.no_grad() |
| def _power_method(self, weight_mat: torch.Tensor, n_power_iterations: int) -> None: |
| # See original note at torch/nn/utils/spectral_norm.py |
| # NB: If `do_power_iteration` is set, the `u` and `v` vectors are |
| # updated in power iteration **in-place**. This is very important |
| # because in `DataParallel` forward, the vectors (being buffers) are |
| # broadcast from the parallelized module to each module replica, |
| # which is a new module object created on the fly. And each replica |
| # runs its own spectral norm power iteration. So simply assigning |
| # the updated vectors to the module this function runs on will cause |
| # the update to be lost forever. And the next time the parallelized |
| # module is replicated, the same randomly initialized vectors are |
| # broadcast and used! |
| # |
| # Therefore, to make the change propagate back, we rely on two |
| # important behaviors (also enforced via tests): |
| # 1. `DataParallel` doesn't clone storage if the broadcast tensor |
| # is already on correct device; and it makes sure that the |
| # parallelized module is already on `device[0]`. |
| # 2. If the out tensor in `out=` kwarg has correct shape, it will |
| # just fill in the values. |
| # Therefore, since the same power iteration is performed on all |
| # devices, simply updating the tensors in-place will make sure that |
| # the module replica on `device[0]` will update the _u vector on the |
| # parallelized module (by shared storage). |
| # |
| # However, after we update `u` and `v` in-place, we need to **clone** |
| # them before using them to normalize the weight. This is to support |
| # backproping through two forward passes, e.g., the common pattern in |
| # GAN training: loss = D(real) - D(fake). Otherwise, engine will |
| # complain that variables needed to do backward for the first forward |
| # (i.e., the `u` and `v` vectors) are changed in the second forward. |
| |
| # Precondition |
| assert weight_mat.ndim > 1 |
| |
| for _ in range(n_power_iterations): |
| # Spectral norm of weight equals to `u^T W v`, where `u` and `v` |
| # are the first left and right singular vectors. |
| # This power iteration produces approximations of `u` and `v`. |
| self._u = F.normalize(torch.mv(weight_mat, self._v), # type: ignore[has-type] |
| dim=0, eps=self.eps, out=self._u) # type: ignore[has-type] |
| self._v = F.normalize(torch.mv(weight_mat.t(), self._u), |
| dim=0, eps=self.eps, out=self._v) # type: ignore[has-type] |
| |
| def forward(self, weight: torch.Tensor) -> torch.Tensor: |
| if weight.ndim == 1: |
| # Faster and more exact path, no need to approximate anything |
| return F.normalize(weight, dim=0, eps=self.eps) |
| else: |
| weight_mat = self._reshape_weight_to_matrix(weight) |
| if self.training: |
| self._power_method(weight_mat, self.n_power_iterations) |
| # See above on why we need to clone |
| u = self._u.clone(memory_format=torch.contiguous_format) |
| v = self._v.clone(memory_format=torch.contiguous_format) |
| # The proper way of computing this should be through F.bilinear, but |
| # it seems to have some efficiency issues: |
| # https://github.com/pytorch/pytorch/issues/58093 |
| sigma = torch.dot(u, torch.mv(weight_mat, v)) |
| return weight / sigma |
| |
| def right_inverse(self, value: torch.Tensor) -> torch.Tensor: |
| # we may want to assert here that the passed value already |
| # satisfies constraints |
| return value |
| |
| |
| def spectral_norm(module: Module, |
| name: str = 'weight', |
| n_power_iterations: int = 1, |
| eps: float = 1e-12, |
| dim: Optional[int] = None) -> Module: |
| r"""Applies spectral normalization to a parameter in the given module. |
| |
| .. math:: |
| \mathbf{W}_{SN} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})}, |
| \sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2} |
| |
| When applied on a vector, it simplifies to |
| |
| .. math:: |
| \mathbf{x}_{SN} = \dfrac{\mathbf{x}}{\|\mathbf{x}\|_2} |
| |
| Spectral normalization stabilizes the training of discriminators (critics) |
| in Generative Adversarial Networks (GANs) by reducing the Lipschitz constant |
| of the model. :math:`\sigma` is approximated performing one iteration of the |
| `power method`_ every time the weight is accessed. If the dimension of the |
| weight tensor is greater than 2, it is reshaped to 2D in power iteration |
| method to get spectral norm. |
| |
| |
| See `Spectral Normalization for Generative Adversarial Networks`_ . |
| |
| .. _`power method`: https://en.wikipedia.org/wiki/Power_iteration |
| .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957 |
| |
| .. note:: |
| This function is implemented using the parametrization functionality |
| in :func:`~torch.nn.utils.parametrize.register_parametrization`. It is a |
| reimplementation of :func:`torch.nn.utils.spectral_norm`. |
| |
| .. note:: |
| When this constraint is registered, the singular vectors associated to the largest |
| singular value are estimated rather than sampled at random. These are then updated |
| performing :attr:`n_power_iterations` of the `power method`_ whenever the tensor |
| is accessed with the module on `training` mode. |
| |
| .. note:: |
| If the `_SpectralNorm` module, i.e., `module.parametrization.weight[idx]`, |
| is in training mode on removal, it will perform another power iteration. |
| If you'd like to avoid this iteration, set the module to eval mode |
| before its removal. |
| |
| Args: |
| module (nn.Module): containing module |
| name (str, optional): name of weight parameter. Default: ``"weight"``. |
| n_power_iterations (int, optional): number of power iterations to |
| calculate spectral norm. Default: ``1``. |
| eps (float, optional): epsilon for numerical stability in |
| calculating norms. Default: ``1e-12``. |
| dim (int, optional): dimension corresponding to number of outputs. |
| Default: ``0``, except for modules that are instances of |
| ConvTranspose{1,2,3}d, when it is ``1`` |
| |
| Returns: |
| The original module with a new parametrization registered to the specified |
| weight |
| |
| Example:: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> snm = spectral_norm(nn.Linear(20, 40)) |
| >>> snm |
| ParametrizedLinear( |
| in_features=20, out_features=40, bias=True |
| (parametrizations): ModuleDict( |
| (weight): ParametrizationList( |
| (0): _SpectralNorm() |
| ) |
| ) |
| ) |
| >>> torch.linalg.matrix_norm(snm.weight, 2) |
| tensor(1.0081, grad_fn=<AmaxBackward0>) |
| """ |
| weight = getattr(module, name, None) |
| if not isinstance(weight, Tensor): |
| raise ValueError( |
| f"Module '{module}' has no parameter or buffer with name '{name}'" |
| ) |
| |
| if dim is None: |
| if isinstance(module, (torch.nn.ConvTranspose1d, |
| torch.nn.ConvTranspose2d, |
| torch.nn.ConvTranspose3d)): |
| dim = 1 |
| else: |
| dim = 0 |
| parametrize.register_parametrization(module, name, _SpectralNorm(weight, n_power_iterations, dim, eps)) |
| return module |
