Repository: pytorch/pytorch

Issue:
Possible memory leak in spectral_norm
I wanted to give a try to spectral_norm function recently included in master through commit ba046331e8cfae7c93a86a5664fcb5c25f9dbee0

I observe that the same model with spectral_norm() around Conv2d & Linear layers keeps its GPU memory increasing until OOM exception.

Removing spectral_norm() call makes everything run smoothly.

- PyTorch: master
- How you installed PyTorch (conda, pip, source): source
- Build command you used (if compiling from source): python setup.py install
- OS: CentOS Linux release 7.4.1708 (Core)
- PyTorch version: master c96f2624a280ae3e2d9195c35150f6aec85f6a02
- Python version: 3.6.5
- CUDA/cuDNN version: 9.1 / 7.1.3
- GPU models and configuration: V100 x 4
- GCC version (if compiling from source): 4.8
- CMake version: 3.5


Discussion:
cc @crcrpar 
I'm sorry to hear that, but I have used my imple. on ver. 0.3.1 w/o such problems.
Of course the imple. which I have used is a little bit different from the one on master branch due to the `Tensor-Variable` marriage.

- ubuntu 16.04
- cuda 9.0
- cudnn 7.0
- how i installed: pip


@mfuntowicz can you try to come up with a self-containing small script to repro the issue?
Please, find the script below. The discriminator network is the same that I'm using in my project. 

On the same machine described at the beginning of the issue, at each epoch the GPU memory increases by ~2Mb (Nvidia-smi visual stats). If you let this script run for a long time, then it gives you out of memory.

Removing the spectral_norm() calls in _ResidualDownSamplingBlock stabilizes memory consumption.

Let me know if I can do anything else.
Morgan

```python
from argparse import ArgumentParser

import torch
import torch.backends.cudnn as cudnn
import torch.cuda as cuda

from torch.nn import Sequential, Conv2d, ReLU, Linear, Module, AvgPool2d, BatchNorm2d
from torch.nn.functional import binary_cross_entropy_with_logits, avg_pool2d
from torch.nn.utils import spectral_norm
from torch.optim import Adam

__author__ = 'Morgan Funtowicz'


class Flatten(Module):
    def forward(self, x):
        x = x.view(x.size(0), -1)
        return x


class _ResidualDownSamplingBlock(Module):

    def __init__(self, n_in, n_out, ksize, stride=1, padding=1):
        super().__init__()
        self._f = Sequential(
            ReLU(),
            spectral_norm(Conv2d(n_in, n_out, ksize, stride, padding)),
            ReLU(True),
            spectral_norm(Conv2d(n_out, n_out, ksize, stride, padding)),
            AvgPool2d(2, 2)
        )
        self._sc = spectral_norm(Conv2d(n_in, n_out, 1, padding=0))

    def forward(self, x):
        return avg_pool2d(self._sc(x), 2, 2) + self._f(x)


if __name__ == '__main__':
    # Ensure Tensor are allocated as FloatTensor
    cudnn.benchmark = True
    torch.set_default_tensor_type('torch.FloatTensor')
    torch.set_default_dtype(torch.float32)

    # Parse provided arguments
    args_parser = ArgumentParser()
    args_parser.add_argument('-d', type=int, default=-1, dest='device', help='Device to use for training (-1 = CPU)')
    args_parser.add_argument('-b', type=int, default=-16, dest='batch', help='Size of the minibatch')

    args = args_parser.parse_args()
    args.gpu = args.device >= 0 and cuda.is_available()
    device = torch.device("cuda:%d" % args.device if args.gpu else "cpu")

    # Define the model & Optimizer
    model = Sequential(
        _ResidualDownSamplingBlock(3, 64, ksize=3),
        _ResidualDownSamplingBlock(64, 64, ksize=3),
        _ResidualDownSamplingBlock(64, 128, ksize=3),
        _ResidualDownSamplingBlock(128, 128, ksize=3),
        _ResidualDownSamplingBlock(128, 128, ksize=3),
        _ResidualDownSamplingBlock(128, 64, ksize=3),
        BatchNorm2d(64), ReLU(True), Flatten(),
        spectral_norm(Linear(256, 1))
    ).to(device)

    opt = Adam(model.parameters())

    # Train
    for epoch in range(20000):
        print('Starting epoch %d' % epoch)

        x, y = torch.randn((args.batch, 3, 128, 128), device=device), torch.rand((args.batch, 1), device=device)
        y_hat = model(x)

        opt.zero_grad()
        loss = binary_cross_entropy_with_logits(y_hat, y)
        loss.backward()
        opt.step()
```
I think I found a couple of issues with the code and I think I even have a fix. The main thing is that it tries to be too clever with the `_u`. I don't think you actually need gradients for `_u`.
Does the code actually work as expected in terms of calculation?
My (limited) understanding of the spectral norm is that you want the gradient to propagate all the way to the original parameter. The [division by sigma (with .data)](https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/spectral_norm.py#L29) looks wrong to me in this respect.

I'll double check that I have a patch in ~4 hours or so.


Resolving patch:
diff --git a/torch/nn/utils/spectral_norm.py b/torch/nn/utils/spectral_norm.py
--- a/torch/nn/utils/spectral_norm.py
+++ b/torch/nn/utils/spectral_norm.py
@@ -14,32 +14,34 @@ def __init__(self, name='weight', n_power_iterations=1, eps=1e-12):
         self.eps = eps
 
     def compute_weight(self, module):
-        weight = module._parameters[self.name + '_org']
-        u = module._buffers[self.name + '_u']
+        weight = getattr(module, self.name + '_org')
+        u = getattr(module, self.name + '_u')
         height = weight.size(0)
         weight_mat = weight.view(height, -1)
-        for _ in range(self.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`.
-            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
+        with torch.no_grad():
+            for _ in range(self.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`.
+                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 = weight / sigma
         return weight, u
 
     def remove(self, module):
         weight = module._parameters[self.name + '_org']
-        del module._parameters[self.name]
-        del module._buffers[self.name + '_u']
-        del module._parameters[self.name + '_org']
+        delattr(module, self.name)
+        delattr(module, self.name + '_u')
+        delattr(module, self.name + '_org')
         module.register_parameter(self.name, weight)
 
     def __call__(self, module, inputs):
         weight, u = self.compute_weight(module)
         setattr(module, self.name, weight)
-        setattr(module, self.name + '_u', u)
+        with torch.no_grad():
+            getattr(module, self.name).copy_(weight)
 
     @staticmethod
     def apply(module, name, n_power_iterations, eps):
@@ -48,7 +50,9 @@ def apply(module, name, n_power_iterations, eps):
         height = weight.size(0)
 
         u = normalize(weight.new_empty(height).normal_(0, 1), dim=0, eps=fn.eps)
+        delattr(module, fn.name)
         module.register_parameter(fn.name + "_org", weight)
+        module.register_buffer(fn.name, weight)
         module.register_buffer(fn.name + "_u", u)
 
         module.register_forward_pre_hook(fn)


Test patch: