dpoulopoulos

strategy = tf.distribute.MirroredStrategy()

BATCH_SIZE = 64
GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync

train_data = tf.data.Dataset(...).batch(GLOBAL_BATCH_SIZE)

with strategy.scope():
  model = tf.keras.Sequential(...)
  model.compile(...)

dpoulopoulos

BATCH_SIZE = 64

train_data = tf.data.Dataset(...).batch(BATCH_SIZE)

model = tf.keras.Sequential(...)
model.compile(...)

model.fit(train_data, epochs=4)

dpoulopoulos

class MyModel(keras.Model):
  def train_step(self, data):
    # Get the data batch
    inputs, targets = data
    # Get the model's weights
    trainable_vars = self.trainable_variables
    # Forward pass
    with tf.GradientTape() as tape:
      # Get the predictions
      preds = self(inputs, training=True)

dpoulopoulos

from PIL import Image
from torchvision.prototype import models as pm


img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

# Step 1: Load a pre-trained model.
# In this step we will load a ResNet architecture.
weights = pm.ResNet50_Weights.ImageNet1K_V1
model = pm.resnet50(weights=weights)

dpoulopoulos

import torch
import torchvision.models as models
import torchvision.transforms as transforms

from PIL import Image


img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

# Step 1: Load a pre-trained model.

dpoulopoulos

imp = package.PackageImporter(path)
loaded_model = imp.load_pickle(package_name, resource_name)

dpoulopoulos

with package.PackageExporter(path) as exp:
    exp.extern("numpy.**")
    exp.intern("models.**")
    exp.save_pickle(package_name, resource_name, model)

dpoulopoulos

with package.PackageExporter(path) as exp:
    exp.intern("models.**")
    exp.save_pickle(package_name, resource_name, model)

dpoulopoulos

from torch import package

path = "/tmp/dcgan.pt"
package_name = "dcgan"
resource_name = "model.pkl"

with package.PackageExporter(path) as exp:
    exp.save_pickle(package_name, resource_name, model)

dpoulopoulos

def run_model(model):
    num_images = 64
    noise, _ = model.buildNoiseData(num_images)
    with torch.no_grad():
      generated_images = model.test(noise)

    # let's plot these images using torchvision and matplotlib
    import matplotlib.pyplot as plt
    import torchvision
    plt.imshow(torchvision.utils.make_grid(generated_images).permute(1, 2, 0).cpu().numpy())