Spectral Normalisation for ConvND Layers

[txir-met]

Hi, I'm trying to do Spectral Norm for 2D and 3D convolutional layers. I found the code for a spectral normaliser in this notebook but it isn't quite working for me. The code I'm trying to run is as follows:

class SNConv2D(snt.Conv2D):
  """2D convolution with spectral normalisation."""

  def __init__(self,sn_eps=0.0001, *args,**kwargs):
    """Constructor."""
    super().__init__(*args,**kwargs)
    self._sn_eps = sn_eps
    self.spectral_normalizer = SpectralNormalizer()

  def __call__(self, inputs, is_training=True):
    self._initialize(inputs)
    self.spectral_normalizer._initialize(self.w)
    normed_w = self.spectral_normalizer(self.w, is_training=is_training)
    outputs = tf.matmul(inputs, normed_w)
    if self.with_bias:
      outputs = tf.add(outputs, self.b)
    return outputs

The spectral normaliser class is taken directly from the jupyter notebook. When I try to run this, the tf.matmul function throws the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: In[0] and In[1] must have compatible batch dimensions: >>[16,32,32,4] vs. [3,3,4,4] [Op:BatchMatMulV2]

[16,32,32,4] is the dimension of the input (which is correct) but [3,3,4,4], the output from the normaliser, doesn't seem right. Any help would be appreciated!

Edit: I think I see the issue. The code in the notebooks was written with the idea it would take 2D input (batch * length), like noise in a basic GAN, but I am trying to use it for a conditional GAN, where the input is of shape (batch * length * width * channels). How would one go about using spectral norm in that situation and aligning the shapes?

[txir-met]

Would it work to simply update self.w with the newly normalised weights and then pass the input through the layer? And does the weight matrix need to be reshaped to be 2D before normalisation takes place? I noticed the tfa and pytorch code for spectral norm does this.

[l4fl4m3]

The example in the notebook is performing the spectral norm on the Linear module. For 2D and 3D Conv layers you would need to implement accordingly. You could use the ConvND class and follow the notebook like logic:

class SNConv2D(sonnet.src.conv.ConvND):
  """2D convolution with spectral normalisation."""

  def __init__(self,sn_eps=0.0001, *args,**kwargs):
    """Constructor."""
    super().__init__(num_spatial_dims=2, data_format= "NHWC", *args,**kwargs)
    self._spectral_normalizer = SpectralNormalizer(epsilon=sn_eps)

  def __call__(self, inputs):

    self._initialize(inputs)

    if self.padding_func:
      inputs = tf.pad(inputs, self._padding)
    
    normed_w = self._spectral_normalizer(self.w)

    outputs = tf.nn.convolution(
          inputs,
          normed_w,
          strides=self.stride,
          padding=self.conv_padding,
          dilations=self.rate,
          data_format=self.data_format)
    
    if self.with_bias:
      outputs = tf.nn.bias_add(outputs, self.b, data_format=self.data_format)

    return outputs

And you should be able to do the same for an SNConv3D function, just increment num_spatial_dims to 3 and change data_format to "NDHWC".

[wang6501sfx]

In sonnet ,where is SpectralNormalizer()?