infer demo

.gitattributes

@@ -25,3 +25,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+result/*
+input/*
+output/*
+*.gif
+nc_workspace/*
+flagged/*

README.md

@@ -1,7 +1,7 @@
 ---
 title: BigDL-Nano Inference
 emoji: 🌖
-colorFrom: pink
+colorFrom: blue
 colorTo: pink
 sdk: gradio
 sdk_version: 3.0.13

app.py

@@ -0,0 +1,144 @@
+import gradio as gr
+import numpy as np
+import time
+from data import write_image_tensor, PatchDataModule, prepare_data, image2tensor, tensor2image
+import torch
+from tqdm import tqdm
+from bigdl.nano.pytorch.trainer import Trainer
+from torch.utils.data import DataLoader
+from pathlib import Path
+from torch.utils.data import Dataset
+import datetime
+
+
+device = 'cpu' 
+dtype = torch.float32
+generator = torch.load("models/generator.pt")
+generator.eval()
+generator.to(device, dtype)
+params = {'batch_size': 1,
+          'num_workers': 0}
+
+
+class ImageDataset(Dataset):
+    def __init__(self, img):
+        self.imgs = [image2tensor(img)]
+    def __getitem__(self, idx: int) -> dict:
+        return self.imgs[idx]
+    
+    def __len__(self) -> int:
+        return len(self.imgs)
+
+
+# quantize model
+data_path = Path('data/webcam')
+train_image_dd = prepare_data(data_path)
+dm = PatchDataModule(train_image_dd, patch_size=2**6,
+                     batch_size=2**3, patch_num=2**6)
+train_loader = dm.train_dataloader()
+train_loader_iter = iter(train_loader)
+quantized_model = Trainer.quantize(generator, accelerator=None,
+                                   calib_dataloader=train_loader)
+
+
+def original_transfer(input_img):
+    w, h, _ = input_img.shape
+    print(datetime.datetime.now())
+    print("input size: ", w, h)
+    # resize too large image
+    if w > 3000 or h > 3000:
+        ratio = min(3000 / w, 3000 / h)
+        w = int(w * ratio)
+        h = int(h * ratio)
+    if w % 4 != 0 or h % 4 != 0:
+        NW = int((w // 4) * 4)
+        NH = int((h // 4) * 4)
+        input_img = np.resize(input_img,(NW,NH,3))
+    st = time.perf_counter()
+    dataset = ImageDataset(input_img)
+    loader = DataLoader(dataset, **params)
+    with torch.no_grad():
+        for inputs in tqdm(loader):
+            inputs = inputs.to(device, dtype)
+            st = time.perf_counter()
+            outputs = generator(inputs)
+            ori_time = time.perf_counter() - st
+            ori_time = "{:.3f}s".format(ori_time)
+            ori_image = np.array(tensor2image(outputs[0]))
+            del inputs
+            del outputs
+    return ori_image, ori_time
+
+def nano_transfer(input_img):
+    w, h, _ = input_img.shape
+    print(datetime.datetime.now())
+    print("input size: ", w, h)
+    # resize too large image
+    if w > 3000 or h > 3000:
+        ratio = min(3000 / w, 3000 / h)
+        w = int(w * ratio)
+        h = int(h * ratio)
+    if w % 4 != 0 or h % 4 != 0:
+        NW = int((w // 4) * 4)
+        NH = int((h // 4) * 4)
+        input_img = np.resize(input_img,(NW,NH,3))
+    st = time.perf_counter()
+    dataset = ImageDataset(input_img)
+    loader = DataLoader(dataset, **params)
+    with torch.no_grad():
+        for inputs in tqdm(loader):
+            inputs = inputs.to(device, dtype)
+            st = time.perf_counter()
+            outputs = quantized_model(inputs)
+            nano_time = time.perf_counter() - st
+            nano_time = "{:.3f}s".format(nano_time)
+            nano_image = np.array(tensor2image(outputs[0]))
+            del inputs
+            del outputs
+    return nano_image, nano_time
+
+
+def clear():
+    return None, None, None, None
+    
+
+demo = gr.Blocks()
+
+with demo:
+    gr.Markdown("<h1><center>BigDL-Nano inference demo</center></h1>")
+    with gr.Row().style(equal_height=False):
+        with gr.Column():
+            gr.Markdown('''
+                <h2>Overview</h2>
+                
+                BigDL-Nano is a library in [BigDL 2.0](https://github.com/intel-analytics/BigDL) that allows the users to transparently accelerate their deep learning pipelines (including data processing, training and inference) by automatically integrating optimized libraries, best-known configurations, and software optimizations. </p>
+                
+                The video on the right shows how the user can easily enable quantization using BigDL-Nano (with just a couple of lines of code); you may refer to our [CVPR 2022 demo paper](https://arxiv.org/abs/2204.01715) for more details.
+                ''')
+        with gr.Column():
+            gr.Video(value="nano_quantize_api.mp4")
+    gr.Markdown('''
+            <h2>Demo</h2>
+            
+            This section uses an image stylization example to demostrate the speedup of the above code when using quantization in BigDL-Nano (about 2~3x inference time speedup). The demo is adapted from the original [FSPBT-Image-Translation code](https://github.com/rnwzd/FSPBT-Image-Translation/blob/master/eval.py).
+            ''')
+    with gr.Row().style(equal_height=False):
+        input_img = gr.Image(label="input image", value="Marvelous_Maisel.jpg", source="upload")
+        with gr.Column():
+            ori_but = gr.Button("Standard PyTorch Lightning")
+            nano_but = gr.Button("BigDL-Nano")
+            clear_but = gr.Button("Clear Output")
+    with gr.Row().style(equal_height=False):
+        with gr.Column():
+            ori_time = gr.Text(label="Standard PyTorch Lightning latency")
+            ori_image = gr.Image(label="Standard PyTorch Lightning output image")
+        with gr.Column():
+            nano_time = gr.Text(label="BigDL-Nano latency")
+            nano_image = gr.Image(label="BigDL-Nano output image")
+    
+    ori_but.click(original_transfer, inputs=input_img, outputs=[ori_image, ori_time])
+    nano_but.click(nano_transfer, inputs=input_img, outputs=[nano_image, nano_time])
+    clear_but.click(clear, inputs=None, outputs=[ori_image, ori_time, nano_image, nano_time])
+    
+
+demo.launch(share=True, enable_queue=True)

data.py

@@ -0,0 +1,205 @@
+from typing import Callable, Dict
+import torch
+from torch.utils.data import Dataset
+import torchvision.transforms.functional as F
+from bigdl.nano.pytorch.vision.transforms import transforms
+import pytorch_lightning as pl
+from collections.abc import Iterable
+# image reader writer
+from pathlib import Path
+from PIL import Image
+from typing import Tuple
+
+
+def read_image(filepath: Path, mode: str = None) -> Image:
+    with open(filepath, 'rb') as file:
+        image = Image.open(file)
+        return image.convert(mode)
+
+
+image2tensor = transforms.ToTensor()
+tensor2image = transforms.ToPILImage()
+
+
+def write_image(image: Image, filepath: Path):
+    filepath.parent.mkdir(parents=True, exist_ok=True)
+    image.save(str(filepath))
+
+
+def read_image_tensor(filepath: Path, mode: str = 'RGB') -> torch.Tensor:
+    return image2tensor(read_image(filepath, mode))
+
+
+def write_image_tensor(input: torch.Tensor, filepath: Path):
+    write_image(tensor2image(input), filepath)
+
+
+def get_valid_indices(H: int, W: int, patch_size: int, random_overlap: int = 0):
+
+    vih = torch.arange(random_overlap, H-patch_size -
+                       random_overlap+1, patch_size)
+    viw = torch.arange(random_overlap, W-patch_size -
+                       random_overlap+1, patch_size)
+    if random_overlap > 0:
+        rih = torch.randint_like(vih, -random_overlap, random_overlap)
+        riw = torch.randint_like(viw, -random_overlap, random_overlap)
+        vih += rih
+        viw += riw
+    vi = torch.stack(torch.meshgrid(vih, viw)).view(2, -1).t()
+    return vi
+
+
+def cut_patches(input: torch.Tensor, indices: Tuple[Tuple[int, int]], patch_size: int, padding: int = 0):
+    # TODO use slices to get all patches at the same time ?
+
+    patches_l = []
+    for n in range(len(indices)):
+
+        patch = F.crop(input, *(indices[n]-padding),
+                       *(patch_size+padding*2,)*2)
+        patches_l.append(patch)
+    patches = torch.cat(patches_l, dim=0)
+
+    return patches
+
+
+def prepare_data(data_path: Path, read_func: Callable = read_image_tensor) -> Dict:
+    """
+    Takes a data_path of a folder which contains subfolders with input, target, etc.
+    lablelled by the same names.
+
+    :param data_path: Path of the folder containing data
+    :param read_func: function that reads data and returns a tensor
+    """
+    data_dict = {}
+
+    subdir_names = ["target", "input", "mask"]  # ,"helper"
+
+    # checks only files for which there is an target
+    # TODO check for images
+    name_ls = [file.name for file in (
+        data_path / "target").iterdir() if file.is_file()]  # 数据集大小=3
+    subdirs = [data_path / sdn for sdn in subdir_names]
+    for sd in subdirs:
+        if sd.is_dir():
+            data_ls = []
+            files = [sd / name for name in name_ls]
+            for file in files:
+                tensor = read_func(file)
+                H, W = tensor.shape[-2:]
+                data_ls.append(tensor)
+            # TODO check that all sizes match
+            data_dict[sd.name] = torch.stack(data_ls, dim=0)
+
+    data_dict['name'] = name_ls
+    data_dict['len'] = len(data_dict['name'])
+    data_dict['H'] = H
+    data_dict['W'] = W
+    return data_dict
+
+
+# TODO an image is loaded whenever a patch is needed, this may be a bottleneck
+class DataDictLoader():
+    def __init__(self, data_dict: Dict,
+                 batch_size: int = 16,
+                 max_length: int = 128,
+                 shuffle: bool = False):
+        """
+
+        """
+
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+
+        self.batch_size = batch_size
+
+        self.data_dict = data_dict
+        self.dataset_len = data_dict['len']  # train: 93
+        self.len = self.dataset_len if max_length is None else min(
+            self.dataset_len, max_length)
+        # Calculate # batches
+        num_batches, remainder = divmod(self.len, self.batch_size)
+        if remainder > 0:
+            num_batches += 1
+        self.num_batches = num_batches
+    
+
+    def __iter__(self):
+        if self.shuffle:
+            r = torch.randperm(self.dataset_len)
+            self.data_dict = {k: v[r] if isinstance(
+                v, Iterable) else v for k, v in self.data_dict.items()}
+        self.i = 0
+        return self
+
+    def __next__(self):
+        if self.i >= self.len:
+            raise StopIteration
+        batch = {k: v[self.i:self.i+self.batch_size]
+                 if isinstance(v, Iterable) else v for k, v in self.data_dict.items()}
+
+        self.i += self.batch_size
+        return batch
+
+    def __len__(self):
+        return self.num_batches
+
+
+class PatchDataModule(pl.LightningDataModule):
+
+    def __init__(self, data_dict,
+                 patch_size: int = 2**5,
+                 batch_size: int = 2**4,
+                 patch_num: int = 2**6):
+        super().__init__()
+        self.data_dict = data_dict
+        self.H, self.W = data_dict['H'], data_dict['W']
+        self.len = data_dict['len']
+
+        self.batch_size = batch_size
+        self.patch_size = patch_size  # 32
+        self.patch_num = patch_num  # 64
+
+    def dataloader(self, data_dict,  **kwargs):
+        return DataDictLoader(data_dict, **kwargs)
+
+    def train_dataloader(self):
+        patches = self.cut_patches()
+        return self.dataloader(patches, batch_size=self.batch_size, shuffle=True,
+                               max_length=self.patch_num)  # patch num = 64
+
+    def val_dataloader(self):
+        return self.dataloader(self.data_dict, batch_size=1)
+
+    def test_dataloader(self):
+        return self.dataloader(self.data_dict)  # TODO batch size
+
+    def cut_patches(self):
+        # TODO cycle once
+        patch_indices = get_valid_indices(
+            self.H, self.W, self.patch_size, self.patch_size//4)
+        dd = {k: cut_patches(
+            v, patch_indices, self.patch_size) for k, v in self.data_dict.items()
+            if isinstance(v, torch.Tensor)
+        }
+        threshold = 0.1
+        mask_p = torch.mean(
+            dd.get('mask', torch.ones_like(dd['input'])), dim=(-1, -2, -3))
+        masked_idx = (mask_p > threshold).nonzero(as_tuple=True)[0]
+        dd = {k: v[masked_idx] for k, v in dd.items()}
+        dd['len'] = len(masked_idx)
+        dd['H'], dd['W'] = (self.patch_size,)*2
+
+        return dd
+
+
+class ImageDataset(Dataset):
+    def __init__(self, file_paths: Iterable, read_func: Callable = read_image_tensor):
+        self.file_paths = file_paths
+
+    def __getitem__(self, idx: int) -> dict:
+        file = self.file_paths[idx]
+        return read_image_tensor(file), file.name
+
+    def __len__(self) -> int:
+        return len(self.file_paths)