CODEWITHSUNDEEP
pythontensorflowkerasdeep-learning
2nebin
2nebin

Reputation: 164

Are there easier ways to excute model's layers in Tensorflow 2.0?

Let's assume that I make a custom model by using Keras subclassing API in Tensorflow 2.0 like below.

class MyModel(Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.conv1 = Conv2D(32, 3, activation='relu')
    self.flatten = Flatten()
    self.d1 = Dense(128, activation='relu')
    self.d2 = Dense(10, activation='softmax')

  def call(self, x):
    x = self.conv1(x)
    x = self.flatten(x)
    x = self.d1(x)
    return self.d2(x)

model = MyModel()

If there aren't many layers, it's not that annoying to write 'def call' of the class.

All I need to do is just writing down 3~4 lines of layers like above.

But what if the number of layer increases like 50~60 or even higher?

Do I have to type all of these layers? Wouldn't there be a much better way I don't know?

I look forward to your answer. Thank you in advance!

Upvotes: 0

Views: 76

Answers (2)

Nicolas Gervais
Nicolas Gervais

Reputation: 36684

You can create a loop to create a number of layers at once, then call them using a for loop.

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from functools import partial

x = tf.random.uniform((5, 2), 0, 1)

class MyModel(tf.keras.Model):
  def __init__(self):
    super(MyModel, self).__init__()
    dense_layer = partial(tf.keras.layers.Dense, units=3, activation='relu')
    self.layer_list = [dense_layer(name=f'layer_{i}') for i in range(3)]
    self.out = tf.keras.layers.Dense(1, activation='sigmoid')

  def call(self, x, training=None, **kwargs):
    for layer in self.layer_list:
        x = layer(x)
    x = self.out(x)
    return x

model = MyModel()
model.build(input_shape=(5, 2))

print(list(map(lambda x: x.name, model.layer_list)))

['layer_0', 'layer_1', 'layer_2']

model(x)

<tf.Tensor: shape=(5, 1), dtype=float32, numpy=
array([[0.51850706],
       [0.5285746 ],
       [0.51396513],
       [0.5171388 ],
       [0.50938624]], dtype=float32)>

Some architectures repeat themselves (e.g., a succession of convolutional layers and max pooling layers multiple times), so you can create function that return a number of layers, and then use a loop to create many of them. To avoid repeating everything, I used functools.partial, which creates a callable with default arguments.

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from functools import partial

x = tf.random.uniform((5, 256, 256, 3), 0, 1)

def conv_block(units):
    conv = partial(tf.keras.layers.Conv2D, kernel_size=(3, 3), activation='relu')
    pool = partial(tf.keras.layers.MaxPooling2D, pool_size=(2, 2))
    layer_dict = {
        'conv1': conv(units),
        'pool1': pool(),
        'conv2': conv(units*2),
        'pool2': pool(),
        'conv3': conv(units*3)
    }
    return layer_dict

class MyModel(tf.keras.Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.block_list = [conv_block(i) for i in range(1, 3)]
    self.flat = tf.keras.layers.Flatten()
    self.out = tf.keras.layers.Dense(1, activation='sigmoid')

  def call(self, x, training=None, **kwargs):
    for block in self.block_list:
        for layer in block.values():
            x = layer(x)
    x = self.flat(x)
    x = self.out(x)
    return x

model = MyModel()
model.build(input_shape=(5, 256, 256, 3))

model(x)

<tf.Tensor: shape=(5, 1), dtype=float32, numpy=
array([[0.48275623],
       [0.48887327],
       [0.49217385],
       [0.48883903],
       [0.48933515]], dtype=float32)>

As you can see, I generated this block twice:

model.layers

[<tensorflow.python.keras.layers.convolutional.Conv2D at 0x1b404f47b08>,
 <tensorflow.python.keras.layers.convolutional.Conv2D at 0x1b40be6cb48>,
 <tensorflow.python.keras.layers.convolutional.Conv2D at 0x1b40d6eab88>,
 <tensorflow.python.keras.layers.pooling.MaxPooling2D at 0x1b40be6cb08>,
 <tensorflow.python.keras.layers.pooling.MaxPooling2D at 0x1b40d6ddc08>,
 <tensorflow.python.keras.layers.convolutional.Conv2D at 0x1b40d6f2188>,
 <tensorflow.python.keras.layers.convolutional.Conv2D at 0x1b40d6f2bc8>,
 <tensorflow.python.keras.layers.convolutional.Conv2D at 0x1b40d6f0648>,
 <tensorflow.python.keras.layers.pooling.MaxPooling2D at 0x1b40d6f2748>,
 <tensorflow.python.keras.layers.pooling.MaxPooling2D at 0x1b40d6f01c8>,
 <tensorflow.python.keras.layers.core.Flatten at 0x1b40d6f8cc8>,
 <tensorflow.python.keras.layers.core.Dense at 0x1b40d6ee248>]

Upvotes: 1

dtlam26
dtlam26

Reputation: 1600

Creating class in TF 2.0 makes it easier to compete with Pytorch and other deep learning framework. However, if your model is Sequential, stacking layers on top of each other, you can still apply it by this example.

For model that need to execute multiple output, parallel computing like SSD or YOLO, writing comprehensively like a class is better or you can use Model directly by passing multiple inputs and multiple outputs to the model (example)

Upvotes: 0

Related Questions