# First time subclassing a model. Need somehelp

**URL:** <https://discuss.ai.google.dev/t/first-time-subclassing-a-model-need-somehelp/26068>\
**Category:** General Discussion\
**Tags:** keras, tf-probability, models, help\_request\
**Created:** [December 10, 2022, 5:00pm UTC](https://discuss.ai.google.dev/t/first-time-subclassing-a-model-need-somehelp/26068 "2022-12-10T17:00:24Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![P11](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/p11/32/10396_2.png) [@P11](https://discuss.ai.google.dev/u/P11)\
**Post date:** [December 10, 2022, 5:00pm UTC](https://discuss.ai.google.dev/t/first-time-subclassing-a-model-need-somehelp/26068/1 "2022-12-10T17:00:24Z")

</div>

Hi guys, i am trying for the first time to use the subclassing method with tensorflow and i am struggling a bit. i was trying to modify the example provided on the tensorflow website. My aim is to change the training step in the future.

```auto
import numpy as np
import tensorflow_probability as tfp
import tensorflow as tf
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
tfk = tf.keras
tfkl = tf.keras.layers
tfd = tfp.distributions
tfpl = tfp.layers

```

```auto
def NLL(y, distr): 
      return -distr.log_prob(y)

```

```auto
def bnn_model():

    inputs= Input(shape=(1,),name="input layer")
    hidden = Dense(50,activation="relu")(inputs)
    hidden = Dropout(0.1)(hidden, training=True)
    hidden = Dense(50,activation="relu")(hidden)
    hidden = Dropout(0.1)(hidden, training=True)
    params_mc = Dense(2,activation="relu")(hidden)
    modello = Model(inputs=inputs, outputs=params_mc)
    return modello

```

```auto
class MyFunModel(tf.keras.Model):

  def __init__ (self, **kwargs):
    super(). __init__ (**kwargs)

 
  def build(self,input_shape):
    def normal_sp(params): 
      return tfd.Normal(loc=params[:,0:1], scale=1e-3 + tf.math.softplus(0.05 * params[:,1:2]))
    self.mlp1 = bnn_model()
    self.prob= tfp.layers.DistributionLambda(normal_sp, name='normal_sp')

 

  def call(self,x):
    two= self.mlp1(x)
    mu= self.prob(two)
    return mu 

  def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True) # Forward pass
            # Compute the loss value
            # (the loss function is configured in `compile()`)
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        # Update metrics (includes the metric that tracks the loss)
        self.compiled_metrics.update_state(y, y_pred)
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

my_funmodel = MyFunModel(name="the_model")

```

but when i try to fit the model

```auto
optimizer = tf.optimizers.Adam(learning_rate=0.0002)

my_funmodel.build((-1,))

my_funmodel.compile(optimizer=optimizer,loss= NLL ) 
history_my_funmodel = my_funmodel.fit(X_train, y_train, epochs=2000, verbose=0, batch_size=batch_size,validation_data=(X_val,y_val) )
plt.plot(history_my_funmodel.history['loss'], label='loss', linewidth=3)
plt.plot(history_my_funmodel.history['val_loss'])
plt.legend()
plt.ylabel('Loss')
plt.xlabel('Epochs')
plt.ylim(top=100, bottom=0)
plt.show()

```

all i get is an empty plot and a fkload of NaN

what can i do to solve this issue? thanks

---

<div class="post-metadata">

**Author:** ![Dennis](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/dennis/32/9676_2.png) [@Dennis](https://discuss.ai.google.dev/u/Dennis)\
**Post date:** [December 11, 2022, 7:21am UTC](https://discuss.ai.google.dev/t/first-time-subclassing-a-model-need-somehelp/26068/2 "2022-12-11T07:21:08Z")

</div>

@P11 You do not need to `compile` and `fit` the model.

Simply call the train\_step while looping:

```auto
for train_step in range(10):
  loss = my_funmodel.train_step([X_train, y_train])
  print(loss)

```

Update:

```auto
def build(self,input_shape):
  ...
  self.optimizer = tf.optimizers.Adam(learning_rate=0.0002)

# my_funmodel.compile(optimizer=optimizer,loss= NLL ) 
# history_my_funmodel = my_funmodel.fit(... )

```

Note:  
You probably need to create your own `self.metrics` and `self.compiled_metrics` function.  
Try to `return loss` only in the first step and comment the metrics lines, to see how the model behaves.