models.py

# -*- coding: utf-8 -*-
"""models.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1ibfKtpxC_hIhZlPbefCoqpAS7jTdyiFw

## Models in TensorFlow

In TensorFlow, you always need to define models to train a machine learning model. A model consists of layers that conduct operations and can be reused in the model's structure. Let's get started.
"""

# Loading necessary libraries
import tensorflow as tf
import numpy as np

"""### Layer

In TensorFlow, we can implement layers using the high-level [tf.Module](https://www.tensorflow.org/api_docs/python/tf/Module) class.
"""

class SampleLayer(tf.Module):
  """
  We define the layer with a class that inherited the structure of tf.Module class.
  """
  def __init__(self, name=None):
    super().__init__(name=name)

    # Define a trainable variable
    self.x = tf.Variable([[1.0, 3.0]], name="x_trainable")

    # Define a non-trainable variable
    self.y = tf.Variable(2.0, trainable=False, name="y_non_trainable")
  def __call__(self, input):
    return self.x * input + self.y

# Initialize the layer
# Here, __call__ function will not be called
simple_layer = SampleLayer(name="my_layer")

# Call the layer and extract some information
output = simple_layer(tf.constant(1.0))
print("Output:", output)
print("Layer name:", simple_layer.name)
print("Trainable variables:", simple_layer.trainable_variables)

"""### Model

Now. let's define a model. A model consists of multiple layers.
"""

class Model(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.layer_1 = SampleLayer('layer_1')
    self.layer_2 = SampleLayer('layer_2')

  def __call__(self, x):
    x = self.layer_1(x)
    output = self.layer_2(x)
    return output

# Initialize the model
custom_model = Model(name="model_name")

# Call the model
# Call the layer and extract some information
output = custom_model(tf.constant(1.0))
print("Output:", output)
print("Model name:", custom_model.name)
print("Trainable variables:", custom_model.trainable_variables)

"""### Keras Models

Keras is a high-level API that is part of TensorFlow now. You can use [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) to define a model. You can also use the collection of [tf.keras.layers](https://www.tensorflow.org/api_docs/python/tf/keras/layers) for your convenience. It's straightforward as below to define a model that has two fully-connected layers:
"""

class CustomModel(tf.keras.Model):

  def __init__(self):
    super(CustomModel, self).__init__()
    self.layer_1 = tf.keras.layers.Dense(16, activation=tf.nn.relu)
    self.layer_2 = tf.keras.layers.Dense(32, activation=None)

  def call(self, inputs):
    x = self.layer_1(inputs)
    out = self.layer_2(inputs)
    return out

# Create model
custom_model = CustomModel()

# Call the model
# Call the layer and extract some information
output = custom_model(tf.constant([[1.0, 2.0, 3.0]]))
print("Output shape:", output.shape)
print("Model name:", custom_model.name)

# Count total trainable variables
total_trainable_var = np.sum([tf.size(var).numpy() for var in custom_model.trainable_variables])
print("Number of trainable variables:", total_trainable_var)