tf.Keras subclass Model having subclass Layer generate AttributeError for save_weights()

Question

I tried this after viewing F. Chollet youtube video "Inside TensorFlow: tf.Keras (part 1). In this session, he discussed how to create custom layer and model by subclassing. I ran this on colab with:

try:
  # %tensorflow_version only exists in Colab.
  %tensorflow_version 2.x
except Exception:
  pass

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.layers import Layer

from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import BinaryCrossentropy

I defined "Linear" layer like this:

class Linear(Layer):

  def __init__(self, units=32, **kwargs):
    super(Linear, self).__init__(**kwargs)
    self.units = units 

  def build(self, input_shape):
    #print("build: input_shape = {}".format(input_shape))
    self.w = self.add_weight(shape=(input_shape[-1], self.units), initializer='random_normal', trainable=True)
    self.b = self.add_weight(shape=(self.units,), initializer='zeros', trainable=True)

  def call(self, inputs):
    return tf.matmul(inputs, self.w) + self.b

  def get_config(self):
    config = super(Linear, self).get_config()
    config.update({'units': self.units})

    return config

and MLP model subclass like this:

class MLP(Model):
  def __init__(self, **kwargs):
    super(MLP, self).__init__(**kwargs)
    self.linear = Linear(1)

  def call(self, inputs):
    return self.linear(inputs)

Note that I have simplified the MLP to have only 1 layer instead of 3 in the video.

I compiled and trained:

x_train = np.random.randn(1024, 64)
y_train = np.random.randint(0, 2, size=(1024, 1))
mlp = MLP()
mlp.compile(optimizer=Adam(), loss=BinaryCrossentropy(from_logits=True))
mlp.fit(x=x_train, y=y_train, epochs=10, verbose=0)
loss = mlp.evaluate(x=x_train, y=y_train, batch_size=1024, verbose=0)

Then saving the weights:

mlp.save_weights('mlp', save_format='tf')

Error:

AttributeError                            Traceback (most recent call last)
<ipython-input-8-00c47f30bead> in <module>()
----> 1 mlp.save_weights('mlp', save_format='tf')
      2 
      3 # export_saved_model(mlp, 'mlp.h5')

8 frames
/tensorflow-2.0.0-rc0/python3.6/tensorflow_core/python/training/tracking/graph_view.py in _escape_local_name(name)
     55   # edges traversed to reach the variable, so we escape forward slashes in
     56   # names.
---> 57   return (name.replace(_ESCAPE_CHAR, _ESCAPE_CHAR + _ESCAPE_CHAR)
     58           .replace(r"/", _ESCAPE_CHAR + "S"))
     59 

AttributeError: 'NoneType' object has no attribute 'replace'

Is there anything I did badly wrong?

Answer 1

Is there anything I did badly wrong?

Not really. When the weights are serialized, not specifying a name for them causes the crash.

Change

    self.w = self.add_weight(shape=(input_shape[-1], self.units), initializer='random_normal', trainable=True)
    self.b = self.add_weight(shape=(self.units,), initializer='zeros', trainable=True)

to

    self.w = self.add_weight(name='w',shape=(input_shape[-1], self.units), initializer='random_normal', trainable=True)
    self.b = self.add_weight(name='b',shape=(self.units,), initializer='zeros', trainable=True)

https://github.com/tensorflow/tensorflow/issues/26811#issuecomment-474255444