Build an Event-Driven Neural Style Transfer Application Using AWS Lambda

Original Source Here

Build an Event-Driven Neural Style Transfer Application Using AWS Lambda

A neural style transfer image (on the right) is generated after 5 epochs and 100 steps per epoch (Image by Author)

To build a production-ready ML application and ensure its stability in the long run, we need to take care of a long checklist of requirements which include the ease with which the models could be iterated, reproducibility, infrastructure, automation, resources, memory, and so on. On top of that, we need a seamless developer experience. How hard could it be?

Flyte can handle the former set of issues because:

  • it’s a workflow automation platform that helps maintain and reproduce pipelines.
  • it provides the control knobs for infrastructure, resources, and memory.

Also, Flyte simplifies the developer experience. In this blog post, we’ll see how by building a neural style transfer application using Flyte and AWS Lambda. We’ll code the end-to-end pipeline and assign the required compute to run the code. Further, we’ll design an event-driven mechanism to trigger the pipeline and output a stylized image when a user uploads an image. From a user perspective, a stylized output image has to be generated upon uploading an image.

Since the application has to be triggered on an event, i.e., image upload, a more suitable choice for Flyte would be to use AWS Lambda. It is serverless and an event-driven compute service. Our neural style transfer application will leverage the “event-driven feature” of AWS Lambda.

Let’s look at how we could stitch the pipeline automation and event-driven service together using Flyte and AWS Lambda.

Figure 1. An overview of the application (Image by Author)

Application Code

Neural style transfer is applying the style of the style image onto the content image. The output image would be a blend of the content and style images.

To get started with the code, first import and configure the dependencies.

Note: This code is an adaption of the Neural style transfer example from the TensorFlow documentation. To run the code, ensure tensorflow, flytekit, and Pillow libraries are installed through pip.

content_layers and style_layers are the layers of the VGG19 model, which we’ll use to build our model, and the tensor_to_image task converts a tensor to an image.

The first step of the model building process is to fetch the image and preprocess it. Define a @task to load the image and limit its maximum dimension to 512 pixels.

The preprocess_img task downloads the content and style image files, and resizes them using the load_img function.

With the data ready to be used by the model, define a VGG19 model that returns the style and content tensors.

class StyleContentModel(tf.keras.models.Model):
def __init__(self, style_layers, content_layers):
super(StyleContentModel, self).__init__()
self.vgg = vgg_layers(style_layers + content_layers)
self.style_layers = style_layers
self.content_layers = content_layers
self.num_style_layers = len(style_layers)
self.vgg.trainable = False
def call(self, inputs):
"Expects float input in [0,1]"
inputs = inputs * 255.0
preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)
outputs = self.vgg(preprocessed_input)
style_outputs, content_outputs = (
outputs[: self.num_style_layers],
outputs[self.num_style_layers :],
style_outputs = [gram_matrix(style_output) for style_output in style_outputs]
content_dict = {
content_name: value
for content_name, value in zip(self.content_layers, content_outputs)
style_dict = {
style_name: value
for style_name, value in zip(self.style_layers, style_outputs)
return {"content": content_dict, "style": style_dict}

The vgg_layers function returns a list of intermediate layer outputs on top of which the model is built (note that we’re using a pretrained VGG network), and the gram_matrix function literally describes the style of an image. When the model is called on an image, it returns the gram matrix of the style_layers and the content of the content_layers.

Next comes the implementation of the style transfer algorithm. Calculate the total loss (style + content) by considering the weighted combination of the two losses.

def style_content_loss(outputs, content_targets, style_targets):
style_outputs = outputs["style"]
content_outputs = outputs["content"]
style_loss = tf.add_n(
tf.reduce_mean((style_outputs[name] - style_targets[name]) ** 2)
for name in style_outputs.keys()
style_loss *= style_weight / len(style_layers)
content_loss = tf.add_n(
tf.reduce_mean((content_outputs[name] - content_targets[name]) ** 2)
for name in content_outputs.keys()
content_loss *= content_weight / len(content_layers)
loss = style_loss + content_loss
return loss

Call style_content_loss from within tf.GradientTape to update the image.

@task(requests=Resources(cpu="1", mem="5Gi", storage="5Gi", ephemeral_storage="5Gi"))
def train_step(
image: tf.Variable, content_image: tf.Tensor, style_image: tf.Tensor
) -> tf.Variable:
opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
extractor = StyleContentModel(style_layers, content_layers)
style_targets = extractor(style_image)["style"]
content_targets = extractor(content_image)["content"]
with tf.GradientTape() as tape:
outputs = extractor(image)
loss = style_content_loss(outputs, content_targets, style_targets)
loss += total_variation_weight * tf.image.total_variation(image)

grad = tape.gradient(loss, image)
opt.apply_gradients([(grad, image)])
return image

The train_step task initializes the style and content target values (tensors), computes the total variation loss, runs gradient descent, applies the processed gradients, and clips the pixel values of the image between 0 and 1. Define the clip_0_1 function as follows:

def clip_0_1(image):
return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

Create a @dynamic workflow to trigger the train_step task for a specified number of epochs and steps_per_epoch.

tf.Variable stores the content image. When it is called from within tf.GradientTape, the image, a tf.Variable is watched and the operations are recorded for automatic differentiation.

Lastly, define a @workflow to encapsulate the tasks and generate a stylized image.

Once the pipeline is deployed, the subsequent step would be to set up the S3 bucket and configure Lambda.

Configure AWS S3 Bucket and Lambda

Images will be uploaded to the S3 bucket, and Lambda will be used to trigger the Flyte workflow as soon as an image is uploaded.

S3 Bucket

To configure the S3 bucket,

  1. Open the Amazon S3 console.
  2. Choose Buckets.
  3. Choose Create bucket.
  4. Give the bucket a name, e.g., “neural-style-transfer”.
  5. Choose the appropriate AWS region (make sure Lambda is created in the same AWS region).
  6. Block or unblock public access (this tutorial assumes that public access is granted).
  7. Choose Create bucket.


A Lambda function can be created from scratch, through a blueprint, a container image, or a serverless app repository. Blueprint can be chosen to fetch sample lambda code, in our case, an S3 blueprint. However, since we need to connect to FlyteRemote from within Lambda, we have to install the flytekit library. Library installation within Lambda is possible through the zip file or container image approach.

Zip file is the easiest approach to get flytekit into Lambda, but due to the size limitations it imposes on the zip file, a much more feasible way would be to use the container image approach.

Container Image

To create a container image on your machine:

1. Create a project directory (e.g., lambda) to accommodate the lambda function.

2. Create 4 files in the directory:, Dockerfile, requirements.txt, and flyte.config.

├── Dockerfile
├── flyte.config
└── requirements.txt

3. encapsulate the code to fetch the uploaded image, instantiate a FlyteRemote object, and trigger the Flyte workflow.

ℹ️ FlyteRemote provides a programmatic interface to interact with the Flyte backend.

Make sure to fill in the endpoint, default_project (e.g. flytesnacks), default_domain (e.g. development), and the name of the launch plan (e.g. neural_style_transfer.example.neural_style_transfer_wf).

4. flyte.config: add configuration to connect to Flyte through FlyteRemote.

5. requirements.txt


6. Dockerfile: copy, flyte.config, and requirements.txt to the root. Instantiate CMD to the handler that is used in the file.

7. Build a Docker image in the project directory using the command:

docker build -t neural-style-transfer .

8. Authenticate Docker CLI to the Amazon ECR registry.

aws ecr get-login-password --region <us-east-1> | docker login --username AWS --password-stdin <123456789012>.dkr.ecr.<us-east-1>

Make sure to replace text in <>.

9. Create a repository in the ECR.

10. Tag your Docker image and push the image to the newly-created repository.

docker tag neural-style-transfer:latest <123456789012>.dkr.ecr.<us-east-1> push <123456789012>.dkr.ecr.<us-east-1>

Make sure to replace text in <> in the registry details.

That’s it! You now have your image in the ECR.

Lambda Configuration

To configure Lambda,

  1. Open the Functions page of the Lambda console.
  2. Choose Create function.
  3. Choose Container image.
  4. Enter the function name (e.g., s3-lambda).
  5. Give the Container Image URI (should be available in Amazon ECR console -> Repositories dashboard).
  6. Choose Create function.

You now have the lambda configured!


S3 bucket and Lambda are currently separate entities. To trigger Lambda as soon as an image is uploaded to the S3 bucket, we must establish a connection between them.

Connecting them also requires setting up the required permissions. But before configuring the permissions, copy the bucket and Lambda ARNs.

Bucket ARN:

  1. Open the Amazon S3 console.
  2. Choose Buckets.
  3. Choose your bucket.
  4. Choose Properties.
  5. Copy the ARN.

Lambda ARN:

  1. Open the Functions page of the Lambda console.
  2. Choose Functions.
  3. Choose your Lambda.
  4. Choose Configuration and then choose Permissions.
  5. Click on the role in Execution role.
  6. Copy the ARN.

S3 Bucket

To set up permissions for the S3 bucket:

  1. Go to the S3 bucket you created.
  2. Select Permissions.
  3. In the Bucket policy, choose Edit.
  4. Add the following policy:

Make sure to fill in the Lambda execution role ARN and the S3 bucket ARN.


To set up permissions for the Lambda:

1. Follow steps 1- 4 outlined in the Lambda ARN section.

2. Under Permissions, choose Add Permissions.

3. In the dropdown, choose Create inline policy.

4. Under the JSON tab, paste the following:

Make sure to fill in the S3 bucket ARN.

5. Choose Review policy.

6. For Name, enter a name for your policy.

7. Choose Create policy.

You can add FLYTE_CREDENTIALS_CLIENT_SECRET to the lambda’s environment variables as part of initializing FlyteRemote. To do so:

  1. Follow steps 1–3 outlined in the Lambda ARN section.
  2. Choose Configuration and then choose Environment Variables.
  3. Set the key as FLYTE_CREDENTIALS_CLIENT_SECRET, and the value should be your secret.

Now comes the fun part — linking lambda to the S3 bucket!


To set up the trigger:

  1. Follow steps 1–3 outlined in the Lambda ARN section.
  2. Choose Configuration and then choose Triggers.
  3. Click Add trigger.
  4. In the Select a trigger dropdown, choose S3.
  5. Choose your S3 bucket under Bucket.
  6. Choose Add.
Figure 2. There should be a link established between S3 and Lambda. (Image by Author)

Test the Application

To test the application, upload an image to the S3 bucket. On your Flyte console, under the neural style transfer workflow, check if the execution got triggered. The output of the execution should be your stylized image!

Next Steps

To summarize, we’ve built an event-driven application that triggers and executes an ML pipeline on the fly whenever there’s new data. It’s quite easy to productionize the pipeline with Flyte and AWS Lambda, as seen in this tutorial. We can also have a front-end application on top of this flow to make the application even more accessible.

If you want to give feedback about this tutorial or have questions regarding the implementation, please post in the comments!


Trending AI/ML Article Identified & Digested via Granola by Ramsey Elbasheer; a Machine-Driven RSS Bot

%d bloggers like this: