Tutorial 5: Model¶

Overview¶

In this tutorial we will cover:

• Instantiating and Compiling a Model
• The Model Function
  • Custom Models
  • FastEstimator Models
  • Pre-Trained Models
• The Optimizer Function
• Loading Model Weights
• Specifying a Model Name
• Related Apphub Examples

Instantiating and Compiling a model¶

We need to specify two things to instantiate and compile a model:

• model_fn
• optimizer_fn

Model definitions can be implemented in Tensorflow or Pytorch and instantiated by calling fe.build which constructs a model instance and associates it with the specified optimizer.

Model Function¶

model_fn should be a function/lambda function which returns either a tf.keras.Model or torch.nn.Module. FastEstimator provides several ways to specify the model architecture:

• Custom model architecture
• Importing a pre-built model architecture from FastEstimator
• Importing pre-trained models/architectures from PyTorch or TensorFlow

Custom model architecture¶

Let's create a custom model in TensorFlow and PyTorch for demonstration.

In [1]:

# Some preliminary imports

import tensorflow as tf

# Since we will be mixing TF and Torch in the tutorial, we need to stop TF from taking all of the GPU memory.
# Normally you would pick either TF or Torch, so you don't need to worry about this.
physical_devices = tf.config.list_physical_devices('GPU')
for device in physical_devices:
    try:
        tf.config.experimental.set_memory_growth(device, True)
    except:
        pass

import torch
import torch.nn as nn

import fastestimator as fe

# Some preliminary imports import tensorflow as tf # Since we will be mixing TF and Torch in the tutorial, we need to stop TF from taking all of the GPU memory. # Normally you would pick either TF or Torch, so you don't need to worry about this. physical_devices = tf.config.list_physical_devices('GPU') for device in physical_devices: try: tf.config.experimental.set_memory_growth(device, True) except: pass import torch import torch.nn as nn import fastestimator as fe

tf.keras.Model¶

In [2]:

def my_model_tf(input_shape=(30, ), num_classes=2):
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Dense(32, activation="relu", input_shape=input_shape))
    model.add(tf.keras.layers.Dense(8, activation="relu"))
    model.add(tf.keras.layers.Dense(num_classes, activation="softmax"))
    return model

model_tf = fe.build(model_fn=my_model_tf, optimizer_fn="adam")

def my_model_tf(input_shape=(30, ), num_classes=2): model = tf.keras.Sequential() model.add(tf.keras.layers.Dense(32, activation="relu", input_shape=input_shape)) model.add(tf.keras.layers.Dense(8, activation="relu")) model.add(tf.keras.layers.Dense(num_classes, activation="softmax")) return model model_tf = fe.build(model_fn=my_model_tf, optimizer_fn="adam")

2022-04-28 17:35:58.413575: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-04-28 17:35:58.978109: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 34777 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB, pci bus id: 0000:bd:00.0, compute capability: 8.0

torch.nn.Module¶

In [3]:

class my_model_torch(nn.Module):
    def __init__(self, num_inputs=30, num_classes=2):
        super().__init__()
        self.layers = nn.Sequential(nn.Linear(num_inputs, 32), 
                                    nn.ReLU(inplace=True), 
                                    nn.Linear(32, 8), 
                                    nn.ReLU(inplace=True),
                                    nn.Linear(8, num_classes))

    def forward(self, x):
        x = self.layers(x)
        x_label = torch.softmax(x, dim=-1)
        return x_label

    
model_torch = fe.build(model_fn=my_model_torch, optimizer_fn="adam")

class my_model_torch(nn.Module): def __init__(self, num_inputs=30, num_classes=2): super().__init__() self.layers = nn.Sequential(nn.Linear(num_inputs, 32), nn.ReLU(inplace=True), nn.Linear(32, 8), nn.ReLU(inplace=True), nn.Linear(8, num_classes)) def forward(self, x): x = self.layers(x) x_label = torch.softmax(x, dim=-1) return x_label model_torch = fe.build(model_fn=my_model_torch, optimizer_fn="adam")

Importing model architecture from FastEstimator¶

Below we import a PyTorch LeNet architecture from FastEstimator. See our Architectures folder for a full list of the architectures provided by FastEstimator.

In [4]:

from fastestimator.architecture.pytorch import LeNet
# from fastestimator.architecture.tensorflow import LeNet  # One can also use a TensorFlow model

model = fe.build(model_fn=LeNet, optimizer_fn="adam")

from fastestimator.architecture.pytorch import LeNet # from fastestimator.architecture.tensorflow import LeNet # One can also use a TensorFlow model model = fe.build(model_fn=LeNet, optimizer_fn="adam")

Importing pre-trained models/architectures from PyTorch or TensorFlow¶

Below we show how to define a model function using a pre-trained resnet model provided by TensorFlow and PyTorch respectively. We load the pre-trained models using a lambda function.

Pre-trained model from tf.keras.applications¶

In [5]:

resnet50_tf = fe.build(model_fn=lambda: tf.keras.applications.ResNet50(weights='imagenet'), optimizer_fn="adam")

resnet50_tf = fe.build(model_fn=lambda: tf.keras.applications.ResNet50(weights='imagenet'), optimizer_fn="adam")

Pre-trained model from torchvision¶

In [6]:

from torchvision import models

resnet50_torch = fe.build(model_fn=lambda: models.resnet50(pretrained=True), optimizer_fn="adam")

from torchvision import models resnet50_torch = fe.build(model_fn=lambda: models.resnet50(pretrained=True), optimizer_fn="adam")

Optimizer function¶

optimizer_fn can be a string or lambda function.

Optimizer from String¶

Specifying a string for the optimizer_fn loads the optimizer with default parameters. The optimizer strings accepted by FastEstimator are as follows:

• Adadelta: 'adadelta'
• Adagrad: 'adagrad'
• Adam: 'adam'
• Adamax: 'adamax'
• RMSprop: 'rmsprop'
• SGD: 'sgd'

Optimizer from Function¶

To specify specific values for the optimizer learning rate or other parameters, we need to pass a lambda function to the optimizer_fn.

In [7]:

# TensorFlow 
model_tf = fe.build(model_fn=my_model_tf, optimizer_fn=lambda: tf.optimizers.Adam(1e-4))

# PyTorch
model_torch = fe.build(model_fn=my_model_torch, optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=1e-4))

# TensorFlow model_tf = fe.build(model_fn=my_model_tf, optimizer_fn=lambda: tf.optimizers.Adam(1e-4)) # PyTorch model_torch = fe.build(model_fn=my_model_torch, optimizer_fn=lambda x: torch.optim.Adam(params=x, lr=1e-4))

If a model function returns multiple models, a list of optimizers can be provided. See the pggan apphub for an example with multiple models and optimizers.

Loading model weights¶

We often need to load the weights of a saved model. Model weights can be loaded by specifying the path of the saved weights using the weights_path parameter. Let's use the resnet models created earlier to showcase this.

Saving model weights¶

Here, we create a temporary directory and use FastEstimator backend to save the weights of our previously created resnet50 models:

In [8]:

import os
import tempfile

model_dir = tempfile.mkdtemp()

# TensorFlow
fe.backend.save_model(resnet50_tf, save_dir=model_dir, model_name= "resnet50_tf")

# PyTorch
fe.backend.save_model(resnet50_torch, save_dir=model_dir, model_name= "resnet50_torch")

import os import tempfile model_dir = tempfile.mkdtemp() # TensorFlow fe.backend.save_model(resnet50_tf, save_dir=model_dir, model_name= "resnet50_tf") # PyTorch fe.backend.save_model(resnet50_torch, save_dir=model_dir, model_name= "resnet50_torch")

Out[8]:

'/tmp/tmpfnjigvpx/resnet50_torch.pt'

Loading weights for TensorFlow and PyTorch models¶

In [9]:

# TensorFlow
resnet50_tf = fe.build(model_fn=lambda: tf.keras.applications.ResNet50(weights=None), 
                       optimizer_fn="adam", 
                       weights_path=os.path.join(model_dir, "resnet50_tf.h5"))

# TensorFlow resnet50_tf = fe.build(model_fn=lambda: tf.keras.applications.ResNet50(weights=None), optimizer_fn="adam", weights_path=os.path.join(model_dir, "resnet50_tf.h5"))

In [10]:

# PyTorch
resnet50_torch = fe.build(model_fn=lambda: models.resnet50(pretrained=False), 
                          optimizer_fn="adam", 
                          weights_path=os.path.join(model_dir, "resnet50_torch.pt"))

# PyTorch resnet50_torch = fe.build(model_fn=lambda: models.resnet50(pretrained=False), optimizer_fn="adam", weights_path=os.path.join(model_dir, "resnet50_torch.pt"))

Specifying a Model Name¶

The name of a model can be specified using the model_name parameter. The name of the model is helpful in distinguishing models when multiple are present.

In [11]:

model = fe.build(model_fn=LeNet, optimizer_fn="adam", model_name="LeNet")
print("Model Name: ", model.model_name)

model = fe.build(model_fn=LeNet, optimizer_fn="adam", model_name="LeNet") print("Model Name: ", model.model_name)

Model Name:  LeNet

If a model function returns multiple models, a list of model_names can be given. See the pggan apphub for an illustration with multiple models and model names.

Apphub Examples¶

You can find some practical examples of the concepts described here in the following FastEstimator Apphubs:

• PG-GAN
• Uncertainty Weighted Loss