Merge pull request #659 from prashantsail/dcgan_fashiongen_example

Author: msaroufim

examples/README.md

@@ -7,6 +7,7 @@
 * [Serving object detection model](#serving-object-detection-model)
 * [Serving image segmentation model](#serving-image-segmentation-model)
 * [Serving huggingface transformers model](#serving-huggingface-transformers)
+* [Serving image generator model](#example-to-serve-GAN-model)
 * [Serving machine translation model](#serving-neural-machine-translation)
 * [Serving waveglow text to speech synthesizer model](#serving-wavegolw-text-to-speech-synthesizer)
 * [Serving multi modal framework model](#Serving-Multi-modal-model)
@@ -85,6 +86,12 @@ The following example demonstrates how to create and serve a pretrained transfor
 
 * [Hugging Face Transformers](Huggingface_Transformers)
 
+## Example to serve GAN model
+
+The following example demonstrates how to create and serve a pretrained DCGAN model from [facebookresearch/pytorch_GAN_zoo](https://github.com/facebookresearch/pytorch_GAN_zoo)
+
+* [GAN Image Generator](dcgan_fashiongen)
+
 ## Serving Neural Machine Translation
 
 The following example demonstrates how to create and serve a neural translation model using fairseq
@@ -113,4 +120,4 @@ The following example demonstrates how to create and serve a complex image class
 
 The following example demonstrates how to create and serve a complex neural machine translation workflow
 
-* [Neural machine Translation workflow](Workflows/nmt_transformers_pipeline)
+* [Neural machine Translation workflow](Workflows/nmt_transformers_pipeline)

examples/dcgan_fashiongen/Readme.md

@@ -0,0 +1,52 @@
+# GAN(Generative Adversarial Networks) models using TorchServe
+- In this example we will demonstrate how to serve a GAN model using TorchServe.
+- We have used a pretrained DCGAN model from [facebookresearch/pytorch_GAN_zoo](https://github.com/facebookresearch/pytorch_GAN_zoo)  
+  (Introduction to [DCGAN on FashionGen](https://pytorch.org/hub/facebookresearch_pytorch-gan-zoo_dcgan/))
+
+### 1. Create a Torch Model Archive
+Execute the following command to create _dcgan_fashiongen.mar_ :
+```
+./create_mar.sh
+```
+The [create_mar.sh](create_mar.sh) script does the following :
+- Download the model's [source code](https://github.com/facebookresearch/pytorch_GAN_zoo/tree/b75dee40918caabb4fe7ec561522717bf096a8cb/models), extract the relevant directory and zip it. (`--extra-files`)
+- Download a checkpoint file [DCGAN_fashionGen-1d67302.pth](https://dl.fbaipublicfiles.com/gan_zoo/DCGAN_fashionGen-1d67302.pth).  (`--serialized-file`)
+- Provide a custom handler - [dcgan_fashiongen_handler.py](dcgan_fashiongen_handler.py). (`--handler`)
+
+Alterantively, you can directly [download the dcgan_fashiongen.mar](https://torchserve.s3.amazonaws.com/mar_files/dcgan_fashiongen.mar)  
+
+### 2. Start TorchServe and Register Model
+```
+mkdir modelstore
+mv dcgan_fashiongen.mar modelstore/
+torchserve --start --ncs --model-store ./modelstore --models dcgan_fashiongen.mar
+```
+
+### 3. Generate Images
+Invoke the predictions API and pass following payload(JSON)
+- **number_of_images** :  Number of images to generate
+- **input_gender** : OPTIONAL; If specified, needs to be one of - `Men`, `Women`
+- **input_category** : OPTIONAL; If specified, needs to be one of - One of - `SHIRTS`, `SWEATERS`, `JEANS`, `PANTS`, `TOPS`, `SUITS & BLAZERS`, `SHORTS`, `JACKETS & COATS`, `SKIRTS`, `JUMPSUITS`, `SWIMWEAR`, `DRESSES`
+- **input_pose** : OPTIONAL; If specified, needs to be one of - `id_gridfs_1`, `id_gridfs_2`, `id_gridfs_3`, `id_gridfs_4`
+
+#### Example
+1. **Create a single image (random gender, category, pose)**
+   ```
+   curl -X POST -d '{"number_of_images":1}' -H "Content-Type: application/json" http://localhost:8080/predictions/dcgan_fashiongen -o img1.jpg
+   ```
+   > Result image should be similar to the one below -  
+   > ![Sample Image 1](sample-output/img1.jpg)
+
+2. **Create '64' images of 'Men' wearing 'Shirts' in 'id_gridfs_1' pose**
+   ```
+   curl -X POST -d '{"number_of_images":64, "input_gender":"Men", "input_category":"SHIRTS", "input_pose":"id_gridfs_1"}' -H "Content-Type: application/json" http://localhost:8080/predictions/dcgan_fashiongen -o img2.jpg
+   ```
+   > Result image should be similar to the one below -  
+   > ![Sample Image 2](sample-output/img2.jpg)
+
+3. **Create '32' images of 'Women' wearing 'Dresses' in 'id_gridfs_3' pose**
+   ```
+   curl -X POST -d '{"number_of_images":32, "input_gender":"Women", "input_category":"DRESSES", "input_pose":"id_gridfs_3"}' -H "Content-Type: application/json" http://localhost:8080/predictions/dcgan_fashiongen -o img3.jpg
+   ```
+   > Result image should be similar to the one below -  
+   > ![Sample Image 3](sample-output/img3.jpg)

examples/dcgan_fashiongen/create_mar.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+REPONAME="pytorch_GAN_zoo"
+SHA="b75dee40918caabb4fe7ec561522717bf096a8cb" #master branch as of 27th Aug 2020
+SRCZIP="$SHA.zip"
+SRCDIR="$REPONAME-$SHA"
+MODELSZIP="models.zip"
+MODELSDIR="models"
+CHECKPOINT="DCGAN_fashionGen-1d67302.pth" #The DCGAN pretrained model as of 27th Aug 2020
+CHECKPOINT_RENAMED="DCGAN_fashionGen.pth"
+
+# Clean Up before exit
+function cleanup {
+  rm -rf $SRCZIP $SRCDIR $MODELSZIP $CHECKPOINT_RENAMED $MODELSDIR
+}
+trap cleanup EXIT
+
+# Download and Extract model's source code
+sudo apt-get install zip unzip
+
+wget https://github.com/facebookresearch/pytorch_GAN_zoo/archive/$SRCZIP
+unzip $SRCZIP
+# Get the models directory from the source code and zip it up
+# This will later be used by torchserve for loading the model
+mv $SRCDIR/models .
+zip -r $MODELSZIP $MODELSDIR
+
+# Download checkpoint
+wget https://dl.fbaipublicfiles.com/gan_zoo/$CHECKPOINT -O $CHECKPOINT_RENAMED
+
+# Create *.mar
+torch-model-archiver --model-name dcgan_fashiongen \
+                     --version 1.0 \
+                     --serialized-file $CHECKPOINT_RENAMED \
+                     --handler dcgan_fashiongen_handler.py \
+                     --extra-files $MODELSZIP \
+                     --force

examples/dcgan_fashiongen/dcgan_fashiongen_handler.py

@@ -0,0 +1,86 @@
+import os
+import zipfile
+import torch
+from io import BytesIO
+from torchvision.utils import save_image
+from ts.torch_handler.base_handler import BaseHandler
+
+MODELSZIP = "models.zip"
+CHECKPOINT = "DCGAN_fashionGen.pth"
+
+
+class ModelHandler(BaseHandler):
+
+    def __init__(self):
+        self.initialized = False
+        self.map_location = None
+        self.device = None
+        self.use_gpu = True
+        self.store_avg = True
+        self.dcgan_model = None
+        self.default_number_of_images = 1
+
+    def initialize(self, context):
+        """
+        Extract the models zip; Take the serialized file and load the model
+        """
+        properties = context.system_properties
+        model_dir = properties.get("model_dir")
+        gpu_id = properties.get("gpu_id")
+
+        self.map_location, self.device, self.use_gpu = \
+            ("cuda", torch.device("cuda:"+str(gpu_id)), True) if torch.cuda.is_available() else \
+            ("cpu", torch.device("cpu"), False)
+
+        # If not already extracted, Extract model source code
+        if not os.path.exists(model_dir + "/models"):
+            with zipfile.ZipFile(model_dir + "/" + MODELSZIP, "r") as zip_ref:
+                zip_ref.extractall(model_dir)
+
+        # Load Model
+        from models.DCGAN import DCGAN
+        self.dcgan_model = DCGAN(useGPU=self.use_gpu, storeAVG=self.store_avg)
+        state_dict = torch.load(model_dir + "/" + CHECKPOINT, map_location=self.map_location)
+        self.dcgan_model.load_state_dict(state_dict)
+
+        self.initialized = True
+
+    def preprocess(self, requests):
+        """
+        Build noise data by using "number of images" and other "constraints" provided by the end user.
+        """
+        preprocessed_data = []
+        for req in requests:
+            data = req.get("data") if req.get("data") is not None else req.get("body", {})
+
+            number_of_images = data.get("number_of_images", self.default_number_of_images)
+            labels = {ky: "b'{}'".format(vl) for ky, vl in data.items() if ky not in ["number_of_images"]}
+
+            noise = self.dcgan_model.buildNoiseDataWithConstraints(number_of_images, labels)
+            preprocessed_data.append({
+                "number_of_images": number_of_images,
+                "input": noise
+            })
+        return preprocessed_data
+
+    def inference(self, preprocessed_data, *args, **kwargs):
+        """
+        Take the noise data as an input tensor, pass it to the model and collect the output tensor.
+        """
+        input_batch = torch.cat(tuple(map(lambda d: d["input"], preprocessed_data)), 0)
+        with torch.no_grad():
+            image_tensor = self.dcgan_model.test(input_batch, getAvG=True, toCPU=True)
+        output_batch = torch.split(image_tensor, tuple(map(lambda d: d["number_of_images"], preprocessed_data)))
+        return output_batch
+
+    def postprocess(self, output_batch):
+        """
+        Create an image(jpeg) using the output tensor.
+        """
+        postprocessed_data = []
+        for op in output_batch:
+            fp = BytesIO()
+            save_image(op, fp, format="JPEG")
+            postprocessed_data.append(fp.getvalue())
+            fp.close()
+        return postprocessed_data

examples/dcgan_fashiongen/sample-output/img1.jpg

@@ -0,0 +1,97 @@
+import os
+import subprocess
+import requests
+import test_utils
+
+from shutil import copy
+from PIL import Image
+from io import BytesIO
+
+CURR_FILE_PATH = os.path.dirname(os.path.realpath(__file__))
+REPO_ROOT_DIR = os.path.normpath(os.path.join(CURR_FILE_PATH, "..", ".."))
+MODELSTORE_DIR = os.path.join(REPO_ROOT_DIR, "modelstore")
+DCGAN_EXAMPLE_DIR = os.path.join(REPO_ROOT_DIR, "examples", "dcgan_fashiongen")
+DCGAN_MAR_FILE = os.path.join(DCGAN_EXAMPLE_DIR, "dcgan_fashiongen.mar")
+
+
+os.environ['MKL_THREADING_LAYER'] = 'GNU'
+# Work around for issue - https://github.com/pytorch/pytorch/issues/37377
+
+def setup_module():
+    test_utils.torchserve_cleanup()
+    create_example_mar()
+
+    os.makedirs(MODELSTORE_DIR, exist_ok=True)  # Create modelstore directory
+    copy(DCGAN_MAR_FILE, MODELSTORE_DIR)  # Copy *.mar to modelstore
+
+    test_utils.start_torchserve(model_store=MODELSTORE_DIR)
+    pass
+
+
+def teardown_module():
+    test_utils.torchserve_cleanup()
+
+    # Empty and Remove modelstore directory
+    test_utils.delete_model_store(MODELSTORE_DIR)
+    os.rmdir(MODELSTORE_DIR)
+
+    delete_example_mar()
+    pass
+
+
+def create_example_mar():
+    # Create only if not already present
+    if not os.path.exists(DCGAN_MAR_FILE):
+        create_mar_cmd = "cd " + DCGAN_EXAMPLE_DIR + ";./create_mar.sh"
+        subprocess.check_call(create_mar_cmd, shell=True)
+
+
+def delete_example_mar():
+    try:
+        os.remove(DCGAN_MAR_FILE)
+    except OSError:
+        pass
+
+
+def test_model_archive_creation():
+    # *.mar created in setup phase
+    assert os.path.exists(DCGAN_MAR_FILE), "Failed to create dcgan mar file"
+
+
+def test_model_register_unregister():
+    reg_resp = test_utils.register_model("dcgan_fashiongen", "dcgan_fashiongen.mar")
+    assert reg_resp.status_code == 200, "Model Registration Failed"
+
+    unreg_resp = test_utils.unregister_model("dcgan_fashiongen")
+    assert unreg_resp.status_code == 200, "Model Unregistration Failed"
+
+
+def test_image_generation_without_any_input_constraints():
+    test_utils.register_model("dcgan_fashiongen", "dcgan_fashiongen.mar")
+    input_json = {}
+    response = requests.post(url="http://localhost:8080/predictions/dcgan_fashiongen", json=input_json)
+    fp = BytesIO(response.content)
+    img = Image.open(fp)
+    # Expect a jpeg of dimension 64 x 64, it contains only 1 image by default
+    assert response.status_code == 200, "Image generation failed"
+    assert img.get_format_mimetype() == "image/jpeg", "Generated image is not a jpeg"
+    assert img.size == (64, 64), "Generated image is not of correct dimensions"
+    test_utils.unregister_model("dcgan_fashiongen")
+
+
+def test_image_generation_with_input_constraints():
+    test_utils.register_model("dcgan_fashiongen", "dcgan_fashiongen.mar")
+    input_json = {
+        "number_of_images": 64,
+        "input_gender": "Men",
+        "input_category": "SHIRTS",
+        "input_pose": "id_gridfs_1"
+    }
+    response = requests.post(url="http://localhost:8080/predictions/dcgan_fashiongen", json=input_json)
+    fp = BytesIO(response.content)
+    img = Image.open(fp)
+    # Expect a jpeg of dimension 530 x 530, it contains 64 images
+    assert response.status_code == 200, "Image generation failed"
+    assert img.get_format_mimetype() == "image/jpeg", "Generated image is not a jpeg"
+    assert img.size == (530, 530), "Generated image is not of correct dimensions"
+    test_utils.unregister_model("dcgan_fashiongen")