replace l2normalize with F.normalize

pytorch · ssnl · May 1, 2018 · Apr 25, 2018 · Apr 25, 2018 · Apr 25, 2018
commit f83f867be65fa981efd82cbc50dfcadb0c33ab4d
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -1306,9 +1306,9 @@ def test_spectral_norm_forward(self):
         _weight, _bias, _u = m.weight_org, m.bias, m.weight_u
         _weight_mat = _weight.view(_weight.size(0), -1)
         _v = torch.mv(_weight_mat.t(), _u)
-        _v = _v / (_v.norm() + 1e-12)
+        _v = F.normalize(_v, dim=0, eps=1e-12)
         _u = torch.mv(_weight_mat, _v)
-        _u = _u / (_u.norm() + 1e-12)
+        _u = F.normalize(_u, dim=0, eps=1e-12)
         _weight.data /= torch.dot(_u, torch.matmul(_weight_mat, _v))
         out_hat = torch.nn.functional.linear(input, _weight, _bias)
         expect_out = m(input)

diff --git a/torch/nn/utils/spectral_norm.py b/torch/nn/utils/spectral_norm.py
@@ -2,15 +2,10 @@
 Spectral Normalization from https://arxiv.org/abs/1802.05957
 """
 import torch
+from torch.nn.functional import normalize
 from torch.nn.parameter import Parameter
 
 
-def l2normalize(v, eps=1e-12):
-    """Scale to inputs norm to 1."""
-    denom = v.norm(p=2) + eps
-    return v / denom
-
-
 class SpectralNorm(object):
 
     def __init__(self, name='weight', n_power_iterations=1, eps=1e-12):
@@ -21,11 +16,11 @@ def __init__(self, name='weight', n_power_iterations=1, eps=1e-12):
     def compute_weight(self, module):
         weight = module._parameters[self.name + '_org']
         u = module._buffers[self.name + '_u']
-        height, _cuda = weight.size(0), weight.is_cuda
+        height = weight.size(0)
         weight_mat = weight.view(height, -1)
         for _ in range(self.n_power_iterations):
-            v = l2normalize(torch.matmul(weight_mat.t(), u), self.eps)
-            u = l2normalize(torch.matmul(weight_mat, v), self.eps)
+            v = normalize(torch.matmul(weight_mat.t(), u), dim=0, eps=self.eps)
+            u = normalize(torch.matmul(weight_mat, v), dim=0, eps=self.eps)
 
         sigma = torch.dot(u, torch.matmul(weight_mat, v))
         weight.data /= sigma
@@ -49,7 +44,7 @@ def apply(module, name, n_power_iterations, eps):
         weight = module._parameters[name]
         height = weight.size(0)
 
-        u = l2normalize(weight.data.new(height).normal_(0, 1), fn.eps)
+        u = normalize(weight.data.new(height).normal_(0, 1), dim=0, eps=fn.eps)
         module.register_parameter(fn.name + "_org", weight)
         module.register_buffer(fn.name + "_u", u)