update training code

.vscode/.ropeproject/config.py

@@ -0,0 +1,114 @@
+# The default ``config.py``
+# flake8: noqa
+
+
+def set_prefs(prefs):
+    """This function is called before opening the project"""
+
+    # Specify which files and folders to ignore in the project.
+    # Changes to ignored resources are not added to the history and
+    # VCSs.  Also they are not returned in `Project.get_files()`.
+    # Note that ``?`` and ``*`` match all characters but slashes.
+    # '*.pyc': matches 'test.pyc' and 'pkg/test.pyc'
+    # 'mod*.pyc': matches 'test/mod1.pyc' but not 'mod/1.pyc'
+    # '.svn': matches 'pkg/.svn' and all of its children
+    # 'build/*.o': matches 'build/lib.o' but not 'build/sub/lib.o'
+    # 'build//*.o': matches 'build/lib.o' and 'build/sub/lib.o'
+    prefs['ignored_resources'] = ['*.pyc', '*~', '.ropeproject',
+                                  '.hg', '.svn', '_svn', '.git', '.tox']
+
+    # Specifies which files should be considered python files.  It is
+    # useful when you have scripts inside your project.  Only files
+    # ending with ``.py`` are considered to be python files by
+    # default.
+    # prefs['python_files'] = ['*.py']
+
+    # Custom source folders:  By default rope searches the project
+    # for finding source folders (folders that should be searched
+    # for finding modules).  You can add paths to that list.  Note
+    # that rope guesses project source folders correctly most of the
+    # time; use this if you have any problems.
+    # The folders should be relative to project root and use '/' for
+    # separating folders regardless of the platform rope is running on.
+    # 'src/my_source_folder' for instance.
+    # prefs.add('source_folders', 'src')
+
+    # You can extend python path for looking up modules
+    # prefs.add('python_path', '~/python/')
+
+    # Should rope save object information or not.
+    prefs['save_objectdb'] = True
+    prefs['compress_objectdb'] = False
+
+    # If `True`, rope analyzes each module when it is being saved.
+    prefs['automatic_soa'] = True
+    # The depth of calls to follow in static object analysis
+    prefs['soa_followed_calls'] = 0
+
+    # If `False` when running modules or unit tests "dynamic object
+    # analysis" is turned off.  This makes them much faster.
+    prefs['perform_doa'] = True
+
+    # Rope can check the validity of its object DB when running.
+    prefs['validate_objectdb'] = True
+
+    # How many undos to hold?
+    prefs['max_history_items'] = 32
+
+    # Shows whether to save history across sessions.
+    prefs['save_history'] = True
+    prefs['compress_history'] = False
+
+    # Set the number spaces used for indenting.  According to
+    # :PEP:`8`, it is best to use 4 spaces.  Since most of rope's
+    # unit-tests use 4 spaces it is more reliable, too.
+    prefs['indent_size'] = 4
+
+    # Builtin and c-extension modules that are allowed to be imported
+    # and inspected by rope.
+    prefs['extension_modules'] = []
+
+    # Add all standard c-extensions to extension_modules list.
+    prefs['import_dynload_stdmods'] = True
+
+    # If `True` modules with syntax errors are considered to be empty.
+    # The default value is `False`; When `False` syntax errors raise
+    # `rope.base.exceptions.ModuleSyntaxError` exception.
+    prefs['ignore_syntax_errors'] = False
+
+    # If `True`, rope ignores unresolvable imports.  Otherwise, they
+    # appear in the importing namespace.
+    prefs['ignore_bad_imports'] = False
+
+    # If `True`, rope will insert new module imports as
+    # `from <package> import <module>` by default.
+    prefs['prefer_module_from_imports'] = False
+
+    # If `True`, rope will transform a comma list of imports into
+    # multiple separate import statements when organizing
+    # imports.
+    prefs['split_imports'] = False
+
+    # If `True`, rope will remove all top-level import statements and
+    # reinsert them at the top of the module when making changes.
+    prefs['pull_imports_to_top'] = True
+
+    # If `True`, rope will sort imports alphabetically by module name instead
+    # of alphabetically by import statement, with from imports after normal
+    # imports.
+    prefs['sort_imports_alphabetically'] = False
+
+    # Location of implementation of
+    # rope.base.oi.type_hinting.interfaces.ITypeHintingFactory In general
+    # case, you don't have to change this value, unless you're an rope expert.
+    # Change this value to inject you own implementations of interfaces
+    # listed in module rope.base.oi.type_hinting.providers.interfaces
+    # For example, you can add you own providers for Django Models, or disable
+    # the search type-hinting in a class hierarchy, etc.
+    prefs['type_hinting_factory'] = (
+        'rope.base.oi.type_hinting.factory.default_type_hinting_factory')
+
+
+def project_opened(project):
+    """This function is called after opening the project"""
+    # Do whatever you like here!

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Python: Current File",
+      "type": "python",
+      "request": "launch",
+      "program": "${file}",
+      "console": "integratedTerminal"
+    }
+  ]
+}

README.md

@@ -66,6 +66,18 @@ python test.py --clip_path ./sample_videos/clips/v32 \
 
 Note that we use 216\*384 images for training, which has aspect ratio of 1:2. During inference, we scale the input to this size and then rescale the output back to the original size.
 
+## Train
+
+We also provide training code for reference. The training can be started by running:
+
+```bash
+python --data_root [root of video samples] \
+       --data_root_imagenet [root of image samples] \
+       --gpu_ids [gpu ids] \
+```
+
+We do not provide the full video dataset due to the copyright issue. For image samples, we retrieve semantically similar images from ImageNet using [this repository](https://github.com/hmmlillian/Gray-Image-Retrieval). Still, one can refer to our code to understand the detailed procedure of augmenting the image dataset to mimic the video frames.
+
 ## Comparison with State-of-the-Arts
 
 <figure>
@@ -73,6 +85,7 @@ Note that we use 216\*384 images for training, which has aspect ratio of 1:2. Du
 </figure>
 
 ## More results
+
 <figure>
 <img src='assets/sample2.png' />
 </figure>

bash.sh

@@ -0,0 +1,30 @@
+python test.py --clip_path ../test_videos/frames/0 --ref_path ../test_videos/refs/0 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/1 --ref_path ../test_videos/refs/1 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/2 --ref_path ../test_videos/refs/2 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/3 --ref_path ../test_videos/refs/3 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/4 --ref_path ../test_videos/refs/4 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/5 --ref_path ../test_videos/refs/5 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/6 --ref_path ../test_videos/refs/6 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/7 --ref_path ../test_videos/refs/7 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/8 --ref_path ../test_videos/refs/8 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/9 --ref_path ../test_videos/refs/9 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/10 --ref_path ../test_videos/refs/10 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/11 --ref_path ../test_videos/refs/11 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/12 --ref_path ../test_videos/refs/12 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/13 --ref_path ../test_videos/refs/13 --output_path ../test_videos/output
+
+python test.py --clip_path ../test_videos/frames/14 --ref_path ../test_videos/refs/14 --output_path ../test_videos/output
+

lib/TrainTransforms.py

@@ -0,0 +1,348 @@
+from __future__ import division
+
+import collections
+import numbers
+import random
+
+import torch
+from PIL import Image
+from skimage import color
+
+import lib.functional as F
+
+__all__ = [
+    "Compose",
+    "Concatenate",
+    "ToTensor",
+    "Normalize",
+    "Resize",
+    "Scale",
+    "CenterCrop",
+    "Pad",
+    "RandomCrop",
+    "RandomHorizontalFlip",
+    "RandomVerticalFlip",
+    "RandomResizedCrop",
+    "RandomSizedCrop",
+    "FiveCrop",
+    "TenCrop",
+    "RGB2Lab",
+]
+
+
+def CustomFunc(inputs, func, *args, **kwargs):
+    im_l = func(inputs[0], *args, **kwargs)
+    im_ab = func(inputs[1], *args, **kwargs)
+    warp_ba = func(inputs[2], *args, **kwargs)
+    warp_aba = func(inputs[3], *args, **kwargs)
+    im_gbl_ab = func(inputs[4], *args, **kwargs)
+    bgr_mc_im = func(inputs[5], *args, **kwargs)
+
+    layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im]
+
+    for l in range(5):
+        layer = inputs[6 + l]
+        err_ba = func(layer[0], *args, **kwargs)
+        err_ab = func(layer[1], *args, **kwargs)
+
+        layer_data.append([err_ba, err_ab])
+
+    return layer_data
+
+
+class Compose(object):
+    """Composes several transforms together.
+
+    Args:
+        transforms (list of ``Transform`` objects): list of transforms to compose.
+
+    Example:
+        >>> transforms.Compose([
+        >>>     transforms.CenterCrop(10),
+        >>>     transforms.ToTensor(),
+        >>> ])
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, inputs):
+        for t in self.transforms:
+            inputs = t(inputs)
+        return inputs
+
+
+class Concatenate(object):
+    """
+    Input: [im_l, im_ab, inputs]
+    inputs = [warp_ba_l, warp_ba_ab, warp_aba, err_pm, err_aba]
+
+    Output:[im_l, err_pm, warp_ba, warp_aba, im_ab, err_aba]
+    """
+
+    def __call__(self, inputs):
+        im_l = inputs[0]
+        im_ab = inputs[1]
+        warp_ba = inputs[2]
+        warp_aba = inputs[3]
+        im_glb_ab = inputs[4]
+        bgr_mc_im = inputs[5]
+        bgr_mc_im = bgr_mc_im[[2, 1, 0], ...]
+
+        err_ba = []
+        err_ab = []
+
+        for l in range(5):
+            layer = inputs[6 + l]
+            err_ba.append(layer[0])
+            err_ab.append(layer[1])
+
+        cerr_ba = torch.cat(err_ba, 0)
+        cerr_ab = torch.cat(err_ab, 0)
+
+        return (im_l, cerr_ba, warp_ba, warp_aba, im_glb_ab, bgr_mc_im, im_ab, cerr_ab)
+
+
+class ToTensor(object):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
+
+    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
+    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0].
+    """
+
+    def __call__(self, inputs):
+        """
+        Args:
+            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+
+        Returns:
+            Tensor: Converted image.
+        """
+        return CustomFunc(inputs, F.to_mytensor)
+
+
+class Normalize(object):
+    """Normalize an tensor image with mean and standard deviation.
+    Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform
+    will normalize each channel of the input ``torch.*Tensor`` i.e.
+    ``input[channel] = (input[channel] - mean[channel]) / std[channel]``
+
+    Args:
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channel.
+    """
+
+    def __call__(self, inputs):
+        """
+        Args:
+            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+
+        Returns:
+            Tensor: Normalized Tensor image.
+        """
+
+        im_l = F.normalize(inputs[0], 50, 1)  # [0, 100]
+        im_ab = F.normalize(inputs[1], (0, 0), (1, 1))  # [-100, 100]
+
+        inputs[2][0:1, :, :] = F.normalize(inputs[2][0:1, :, :], 50, 1)
+        inputs[2][1:3, :, :] = F.normalize(inputs[2][1:3, :, :], (0, 0), (1, 1))
+        warp_ba = inputs[2]
+
+        inputs[3][0:1, :, :] = F.normalize(inputs[3][0:1, :, :], 50, 1)
+        inputs[3][1:3, :, :] = F.normalize(inputs[3][1:3, :, :], (0, 0), (1, 1))
+        warp_aba = inputs[3]
+
+        im_gbl_ab = F.normalize(inputs[4], (0, 0), (1, 1))  # [-100, 100]
+
+        bgr_mc_im = F.normalize(inputs[5], (123.68, 116.78, 103.938), (1, 1, 1))
+
+        layer_data = [im_l, im_ab, warp_ba, warp_aba, im_gbl_ab, bgr_mc_im]
+
+        for l in range(5):
+            layer = inputs[6 + l]
+            err_ba = F.normalize(layer[0], 127, 2)  # [0, 255]
+            err_ab = F.normalize(layer[1], 127, 2)  # [0, 255]
+            layer_data.append([err_ba, err_ab])
+
+        return layer_data
+
+
+class Resize(object):
+    """Resize the input PIL Image to the given size.
+
+    Args:
+        size (sequence or int): Desired output size. If size is a sequence like
+            (h, w), output size will be matched to this. If size is an int,
+            smaller edge of the image will be matched to this number.
+            i.e, if height > width, then image will be rescaled to
+            (size * height / width, size)
+        interpolation (int, optional): Desired interpolation. Default is
+            ``PIL.Image.BILINEAR``
+    """
+
+    def __init__(self, size, interpolation=Image.BILINEAR):
+        assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)
+        self.size = size
+        self.interpolation = interpolation
+
+    def __call__(self, inputs):
+        """
+        Args:
+            img (PIL Image): Image to be scaled.
+
+        Returns:
+            PIL Image: Rescaled image.
+        """
+        return CustomFunc(inputs, F.resize, self.size, self.interpolation)
+
+
+class RandomCrop(object):
+    """Crop the given PIL Image at a random location.
+
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+        padding (int or sequence, optional): Optional padding on each border
+            of the image. Default is 0, i.e no padding. If a sequence of length
+            4 is provided, it is used to pad left, top, right, bottom borders
+            respectively.
+    """
+
+    def __init__(self, size, padding=0):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+        self.padding = padding
+
+    @staticmethod
+    def get_params(img, output_size):
+        """Get parameters for ``crop`` for a random crop.
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            output_size (tuple): Expected output size of the crop.
+
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
+        """
+        w, h = img.size
+        th, tw = output_size
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+        return i, j, th, tw
+
+    def __call__(self, inputs):
+        """
+        Args:
+            img (PIL Image): Image to be cropped.
+
+        Returns:
+            PIL Image: Cropped image.
+        """
+        if self.padding > 0:
+            inputs = CustomFunc(inputs, F.pad, self.padding)
+
+        i, j, h, w = self.get_params(inputs[0], self.size)
+        return CustomFunc(inputs, F.crop, i, j, h, w)
+
+
+class CenterCrop(object):
+    """Crop the given PIL Image at a random location.
+
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+        padding (int or sequence, optional): Optional padding on each border
+            of the image. Default is 0, i.e no padding. If a sequence of length
+            4 is provided, it is used to pad left, top, right, bottom borders
+            respectively.
+    """
+
+    def __init__(self, size, padding=0):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+        self.padding = padding
+
+    @staticmethod
+    def get_params(img, output_size):
+        """Get parameters for ``crop`` for a random crop.
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            output_size (tuple): Expected output size of the crop.
+
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
+        """
+        w, h = img.size
+        th, tw = output_size
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = (h - th) // 2
+        j = (w - tw) // 2
+        return i, j, th, tw
+
+    def __call__(self, inputs):
+        """
+        Args:
+            img (PIL Image): Image to be cropped.
+
+        Returns:
+            PIL Image: Cropped image.
+        """
+        if self.padding > 0:
+            inputs = CustomFunc(inputs, F.pad, self.padding)
+
+        i, j, h, w = self.get_params(inputs[0], self.size)
+        return CustomFunc(inputs, F.crop, i, j, h, w)
+
+
+class RandomHorizontalFlip(object):
+    """Horizontally flip the given PIL Image randomly with a probability of 0.5."""
+
+    def __call__(self, inputs):
+        """
+        Args:
+            img (PIL Image): Image to be flipped.
+
+        Returns:
+            PIL Image: Randomly flipped image.
+        """
+
+        if random.random() < 0.5:
+            return CustomFunc(inputs, F.hflip)
+        return inputs
+
+
+class RGB2Lab(object):
+    def __call__(self, inputs):
+        """
+        Args:
+            img (PIL Image): Image to be flipped.
+
+        Returns:
+            PIL Image: Randomly flipped image.
+        """
+
+    def __call__(self, inputs):
+        image_lab = color.rgb2lab(inputs[0])
+        warp_ba_lab = color.rgb2lab(inputs[2])
+        warp_aba_lab = color.rgb2lab(inputs[3])
+        im_gbl_lab = color.rgb2lab(inputs[4])
+
+        inputs[0] = image_lab[:, :, :1]  # l channel
+        inputs[1] = image_lab[:, :, 1:]  # ab channel
+        inputs[2] = warp_ba_lab  # lab channel
+        inputs[3] = warp_aba_lab  # lab channel
+        inputs[4] = im_gbl_lab[:, :, 1:]  # ab channel
+
+        return inputs

lib/videoloader.py

@@ -0,0 +1,301 @@
+import sys
+
+sys.path.insert(0, "..")
+import os
+import random
+
+import cv2
+import numpy as np
+import torch
+import torchvision.utils as vutils
+from PIL import Image
+from skimage import color
+from torch.autograd import Variable
+from utils.flowlib import read_flow
+from utils.util_distortion import CenterPad
+
+import lib.functional as F
+
+cv2.setNumThreads(0)
+
+
+class RGB2Lab(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, inputs):
+        return color.rgb2lab(inputs)
+
+
+class Normalize(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, inputs):
+        inputs[0:1, :, :] = F.normalize(inputs[0:1, :, :], 50, 1)
+        inputs[1:3, :, :] = F.normalize(inputs[1:3, :, :], (0, 0), (1, 1))
+        return inputs
+
+
+class ToTensor(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, inputs):
+        outputs = F.to_mytensor(inputs)  # permute channel and transform to tensor
+        return outputs
+
+
+class RandomErasing(object):
+    def __init__(self, probability=0.6, sl=0.05, sh=0.6):
+        self.probability = probability
+        self.sl = sl
+        self.sh = sh
+
+    def __call__(self, img):
+        img = np.array(img)
+        if random.uniform(0, 1) > self.probability:
+            return Image.fromarray(img)
+
+        area = img.shape[0] * img.shape[1]
+        h0 = img.shape[0]
+        w0 = img.shape[1]
+        channel = img.shape[2]
+
+        h = int(round(random.uniform(self.sl, self.sh) * h0))
+        w = int(round(random.uniform(self.sl, self.sh) * w0))
+
+        if w < img.shape[1] and h < img.shape[0]:
+            x1 = random.randint(0, img.shape[0] - h)
+            y1 = random.randint(0, img.shape[1] - w)
+            img[x1 : x1 + h, y1 : y1 + w, :] = np.random.rand(h, w, channel) * 255
+            return Image.fromarray(img)
+
+        return Image.fromarray(img)
+
+
+class CenterCrop(object):
+    """
+    center crop the numpy array
+    """
+
+    def __init__(self, image_size):
+        self.h0, self.w0 = image_size
+
+    def __call__(self, input_numpy):
+        if input_numpy.ndim == 3:
+            h, w, channel = input_numpy.shape
+            output_numpy = np.zeros((self.h0, self.w0, channel))
+            output_numpy = input_numpy[
+                (h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0, :
+            ]
+        else:
+            h, w = input_numpy.shape
+            output_numpy = np.zeros((self.h0, self.w0))
+            output_numpy = input_numpy[
+                (h - self.h0) // 2 : (h - self.h0) // 2 + self.h0, (w - self.w0) // 2 : (w - self.w0) // 2 + self.w0
+            ]
+        return output_numpy
+
+
+def parse_images(data_root):
+    image_pairs = []
+    subdirs = sorted(os.listdir(data_root))
+    for subdir in subdirs:
+        path = os.path.join(data_root, subdir)
+        if not os.path.isdir(path):
+            continue
+
+        parse_file = os.path.join(path, "pairs_output_new.txt")
+        if os.path.exists(parse_file):
+            with open(parse_file, "r") as f:
+                lines = f.readlines()
+                for line in lines:
+                    line = line.replace("\n", "")
+                    (
+                        image1_name,
+                        image2_name,
+                        reference_video_name,
+                        reference_video_name1,
+                        reference_name1,
+                        reference_name2,
+                        reference_name3,
+                        reference_name4,
+                        reference_name5,
+                        reference_gt1,
+                        reference_gt2,
+                        reference_gt3,
+                    ) = line.split()
+                    image1_name = image1_name.split(".")[0]
+                    image2_name = image2_name.split(".")[0]
+                    reference_video_name = reference_video_name.split(".")[0]
+                    reference_video_name1 = reference_video_name1.split(".")[0]
+                    reference_name1 = reference_name1.split(".")[0]
+                    reference_name2 = reference_name2.split(".")[0]
+                    reference_name3 = reference_name3.split(".")[0]
+                    reference_name4 = reference_name4.split(".")[0]
+                    reference_name5 = reference_name5.split(".")[0]
+
+                    reference_gt1 = reference_gt1.split(".")[0]
+                    reference_gt2 = reference_gt2.split(".")[0]
+                    reference_gt3 = reference_gt3.split(".")[0]
+
+                    flow_forward_name = image1_name + "_forward"
+                    flow_backward_name = image1_name + "_backward"
+                    mask_name = image1_name + "_mask"
+
+                    item = (
+                        image1_name + ".jpg",
+                        image2_name + ".jpg",
+                        reference_video_name + ".jpg",
+                        reference_name1 + ".JPEG",
+                        reference_name2 + ".JPEG",
+                        reference_name3 + ".JPEG",
+                        reference_name4 + ".JPEG",
+                        reference_name5 + ".JPEG",
+                        flow_forward_name + ".flo",
+                        flow_backward_name + ".flo",
+                        mask_name + ".pgm",
+                        reference_gt1 + ".jpg",
+                        reference_gt2 + ".jpg",
+                        reference_gt3 + ".jpg",
+                        path,
+                    )
+                    image_pairs.append(item)
+
+        else:
+            raise (RuntimeError("Error when parsing pair_output_count.txt in subfolders of: " + path + "\n"))
+
+    return image_pairs
+
+
+class VideosDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        data_root,
+        epoch,
+        image_size,
+        image_transform=None,
+        use_google_reference=False,
+        real_reference_probability=1,
+        nonzero_placeholder_probability=0.5,
+    ):
+        self.data_root = data_root
+        self.image_transform = image_transform
+        self.CenterPad = CenterPad(image_size)
+        self.ToTensor = ToTensor()
+        self.CenterCrop = CenterCrop(image_size)
+
+        assert len(self.data_root) > 0, "find no dataroot"
+        self.epoch = epoch
+        self.image_pairs = parse_images(self.data_root)
+        self.real_len = len(self.image_pairs)
+        print("##### parsing image pairs in %s: %d pairs #####" % (data_root, self.real_len))
+        self.image_pairs *= epoch
+        self.use_google_reference = use_google_reference
+        self.real_reference_probability = real_reference_probability
+        self.nonzero_placeholder_probability = nonzero_placeholder_probability
+
+    def __getitem__(self, index):
+        (
+            image1_name,
+            image2_name,
+            reference_video_name,
+            reference_name1,
+            reference_name2,
+            reference_name3,
+            reference_name4,
+            reference_name5,
+            flow_forward_name,
+            flow_backward_name,
+            mask_name,
+            reference_gt1,
+            reference_gt2,
+            reference_gt3,
+            path,
+        ) = self.image_pairs[index]
+        try:
+            I1 = Image.open(os.path.join(path, "input_pad", image1_name))
+            I2 = Image.open(os.path.join(path, "input_pad", image2_name))
+
+            I_reference_video = Image.open(
+                os.path.join(path, "reference_gt", random.choice([reference_gt1, reference_gt2, reference_gt3]))
+            )
+            I_reference_video_real = Image.open(
+                os.path.join(
+                    path,
+                    "reference",
+                    random.choice(
+                        [reference_name1, reference_name2, reference_name3, reference_name4, reference_name5]
+                    ),
+                )
+            )
+
+            flow_forward = read_flow(os.path.join(path, "flow", flow_forward_name))  # numpy
+            flow_backward = read_flow(os.path.join(path, "flow", flow_backward_name))  # numpy
+            mask = Image.open(os.path.join(path, "mask", mask_name))
+
+            # binary mask
+            mask = np.array(mask)
+            mask[mask < 240] = 0
+            mask[mask >= 240] = 1
+
+            # transform
+            I1 = self.image_transform(I1)
+            I2 = self.image_transform(I2)
+            I_reference_video = self.image_transform(self.CenterPad(I_reference_video))
+            I_reference_video_real = self.image_transform(self.CenterPad(I_reference_video_real))
+            flow_forward = self.ToTensor(self.CenterCrop(flow_forward))
+            flow_backward = self.ToTensor(self.CenterCrop(flow_backward))
+            mask = self.ToTensor(self.CenterCrop(mask))
+
+            if np.random.random() < self.real_reference_probability:
+                I_reference_output = I_reference_video_real
+                placeholder = torch.zeros_like(I1)
+                self_ref_flag = torch.zeros_like(I1)
+            else:
+                I_reference_output = I_reference_video
+                placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1)
+                self_ref_flag = torch.ones_like(I1)
+
+            outputs = [
+                I1,
+                I2,
+                I_reference_output,
+                flow_forward,
+                flow_backward,
+                mask,
+                placeholder,
+                self_ref_flag,
+            ]
+
+        except Exception as e:
+            print("problem in, ", path)
+            print(e)
+            return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
+        return outputs
+
+    def __len__(self):
+        return len(self.image_pairs)
+
+
+def batch_lab2rgb_transpose_mc(img_l_mc, img_ab_mc, nrow=8):
+    if isinstance(img_l_mc, Variable):
+        img_l_mc = img_l_mc.data.cpu()
+    if isinstance(img_ab_mc, Variable):
+        img_ab_mc = img_ab_mc.data.cpu()
+
+    if img_l_mc.is_cuda:
+        img_l_mc = img_l_mc.cpu()
+    if img_ab_mc.is_cuda:
+        img_ab_mc = img_ab_mc.cpu()
+
+    assert img_l_mc.dim() == 4 and img_ab_mc.dim() == 4, "only for batch input"
+
+    l_norm, ab_norm = 1.0, 1.0
+    l_mean, ab_mean = 50.0, 0
+    img_l = img_l_mc * l_norm + l_mean
+    img_ab = img_ab_mc * ab_norm + ab_mean
+    pred_lab = torch.cat((img_l, img_ab), dim=1)
+    grid_lab = vutils.make_grid(pred_lab, nrow=nrow).numpy().astype("float64")
+    return (np.clip(color.lab2rgb(grid_lab.transpose((1, 2, 0))), 0, 1) * 255).astype("uint8")

lib/videoloader_imagenet.py

@@ -0,0 +1,290 @@
+import os
+import os.path as osp
+import struct
+
+import cv2
+import numpy as np
+import torch
+import torch.utils.data as data
+import torchvision.transforms as transforms
+from PIL import Image
+from scipy.ndimage import map_coordinates
+from scipy.ndimage.filters import gaussian_filter
+from utils.util_distortion import CenterPadCrop_numpy, Distortion_with_flow, Normalize, RGB2Lab, ToTensor
+
+cv2.setNumThreads(0)
+
+
+def parse_images(dir, with_bad, with_mid):
+    print("dir is: ", dir)
+    dir = osp.expanduser(dir)
+
+    image_pairs = []
+
+    no_mid_cnt, no_bad_cnt = 0, 0
+    for target in sorted(os.listdir(dir)):
+        d = osp.join(dir, target)
+        if not osp.isdir(d):
+            continue
+
+        pair_file = osp.join(d, "pairs.txt")
+
+        if osp.exists(pair_file):
+            with open(pair_file, "r") as f:
+                lines = f.readlines()
+                for line in lines:
+                    pair = line.strip().split(" ")
+                    name0, postfix = pair[0].split(".")
+                    name1, postfix = pair[1].split(".")
+                    if float(pair[2]) > 0:
+                        item0 = (dir, target, name0, name1, 2)
+                        item1 = (dir, target, name1, name0, 2)
+                        image_pairs.append(item0)
+                        image_pairs.append(item1)
+
+        else:
+            raise (RuntimeError("Found no pair.txt in subfolders of: " + d + "\n"))
+
+        if with_mid:
+            pair_file = osp.join(d, "pairs_mid.txt")
+
+            if osp.exists(pair_file):
+                with open(pair_file, "r") as f:
+                    for line in f:
+                        pair = line.strip().split(" ")
+                        name0, postfix = pair[0].split(".")
+                        name1, postfix = pair[1].split(".")
+                        item0 = (dir, target, name0, name1, 0)
+                        item1 = (dir, target, name1, name0, 0)
+                        image_pairs.append(item0)
+                        image_pairs.append(item1)
+
+            else:
+                no_mid_cnt += 1
+
+        if with_bad:
+            pair_file = osp.join(d, "pairs_bad.txt")
+
+            if osp.exists(pair_file):
+                with open(pair_file, "r") as f:
+                    for line in f:
+                        pair = line.strip().split(" ")
+                        name0, postfix = pair[0].split(".")
+                        name1, postfix = pair[1].split(".")
+                        item0 = (dir, target, name0, name1, -1)
+                        item1 = (dir, target, name1, name0, -1)
+                        image_pairs.append(item0)
+                        image_pairs.append(item1)
+
+            else:
+                no_bad_cnt += 1
+
+    if no_bad_cnt > 0:
+        print("find no pairs_bad.txt in %d folders for dir: %s" % (no_bad_cnt, dir))
+
+    if no_mid_cnt > 0:
+        print("find no pairs_mid.txt in %d folders for dir: %s" % (no_mid_cnt, dir))
+
+    return image_pairs
+
+
+def pil_loader(path):
+    with open(path, "rb") as f:
+        with Image.open(f) as img:
+            return img.convert("RGB")
+
+
+def image_loader(path):
+    return pil_loader(path)
+
+
+def combo5_loader(path, real_w, real_h):
+    with open(path, "rb") as f:
+        d = f.read(4)
+        im_sz = struct.unpack("i", d)
+        h = im_sz[0]
+
+        d = f.read(4)
+        im_sz = struct.unpack("i", d)
+        w = im_sz[0]
+
+        d = f.read(4)
+        im_sz = struct.unpack("i", d)
+        d = f.read(im_sz[0])
+        file_bytes = np.asarray(bytearray(d), dtype=np.uint8)
+        img_data_ndarray = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+        img_data_ndarray = cv2.cvtColor(img_data_ndarray, cv2.COLOR_BGR2RGB)
+        warp_ba = Image.fromarray(img_data_ndarray)
+        img_data_ndarray = img_data_ndarray * 0
+        warp_aba = Image.fromarray(img_data_ndarray)
+
+        errs = []
+        empty_array = np.zeros([np.array(warp_ba).shape[0], np.array(warp_ba).shape[1]])
+        for _ in range(5):
+            err_ba = Image.fromarray(empty_array)
+            err_ab = Image.fromarray(empty_array)
+            errs.append([err_ba, err_ab])
+    return errs, warp_ba, warp_aba
+
+
+class VideosDataset_ImageNet(data.Dataset):
+    def __init__(
+        self,
+        data_root,
+        epoch,
+        image_size,
+        with_bad=False,
+        with_mid=False,
+        transforms_imagenet=None,
+        distortion_level=3,
+        brightnessjitter=0,
+        nonzero_placeholder_probability=0.5,
+        extra_reference_transform=None,
+        real_reference_probability=1,
+    ):
+        image_pairs = []
+        curr_image_pairs = parse_images(data_root, with_bad, with_mid)
+        image_pairs += curr_image_pairs
+        print("##### parsing image_a pairs in %s: %d pairs #####" % (data_root, len(curr_image_pairs)))
+        if not image_pairs:
+            raise RuntimeError("Found 0 image_a pairs in all the data_roots")
+
+        self.image_pairs = image_pairs
+        self.transforms_imagenet_raw = transforms_imagenet
+        self.extra_reference_transform = transforms.Compose(extra_reference_transform)
+        self.real_reference_probability = real_reference_probability
+        self.transforms_imagenet = transforms.Compose(transforms_imagenet)
+        self.epoch = epoch
+        self.image_size = image_size
+        self.real_len = len(self.image_pairs)
+        self.image_pairs *= epoch
+        self.distortion_level = distortion_level
+        self.distortion_transform = Distortion_with_flow()
+        self.brightnessjitter = brightnessjitter
+        self.flow_transform = transforms.Compose([CenterPadCrop_numpy(self.image_size), ToTensor()])
+        self.nonzero_placeholder_probability = nonzero_placeholder_probability
+        self.ToTensor = ToTensor()
+        self.Normalize = Normalize()
+
+    def __getitem__(self, index):
+        try:
+            image_id = 0
+            pair_id = index
+
+            combo_path = None
+            image_a_path = None
+            image_b_path = None
+
+            image_names = ["", ""]
+            dir_root, cls_dir, image_names[0], image_names[1], is_good = self.image_pairs[pair_id]
+            sub_dir = osp.join(dir_root, cls_dir)
+            if is_good >= 1:
+                image_a_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[image_id])
+                image_b_path = osp.join(sub_dir, "input", "%s.JPEG" % image_names[1 - image_id])
+            elif is_good == 0:
+                image_a_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[image_id])
+                image_b_path = osp.join(sub_dir, "input_mid", "%s.JPEG" % image_names[1 - image_id])
+            else:
+                image_a_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[image_id])
+                image_b_path = osp.join(sub_dir, "input_bad", "%s.JPEG" % image_names[1 - image_id])
+
+            if np.random.random() > 0.5:
+                image_a_path, image_b_path = image_b_path, image_a_path
+
+            I1 = image_loader(image_a_path)
+            I2 = I1
+            I_reference_video = I1
+            I_reference_video_real = image_loader(image_b_path)
+
+            ## generate the flow
+            height, width = np.array(I2).shape[0], np.array(I2).shape[1]
+            alpha = np.random.rand() * self.distortion_level
+            distortion_range = 50
+            random_state = np.random.RandomState(None)
+            shape = self.image_size[0], self.image_size[1]
+            # dx: flow on the vertical direction; dy: flow on the horizontal direction
+            forward_dx = (
+                gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0)
+                * alpha
+                * 1000
+            )
+            forward_dy = (
+                gaussian_filter((random_state.rand(*shape) * 2 - 1), distortion_range, mode="constant", cval=0)
+                * alpha
+                * 1000
+            )
+
+            for transform in self.transforms_imagenet_raw:
+                if type(transform) is RGB2Lab:
+                    I1_raw = I1
+                I1 = transform(I1)
+            for transform in self.transforms_imagenet_raw:
+                if type(transform) is RGB2Lab:
+                    I2 = self.distortion_transform(I2, forward_dx, forward_dy)
+                    I2_raw = I2
+                I2 = transform(I2)
+            I2[0:1, :, :] = I2[0:1, :, :] + torch.randn(1) * self.brightnessjitter
+
+            I_reference_video = self.extra_reference_transform(I_reference_video)
+            for transform in self.transforms_imagenet_raw:
+                I_reference_video = transform(I_reference_video)
+
+            I_reference_video_real = self.transforms_imagenet(I_reference_video_real)
+
+            flow_forward_raw = np.stack((forward_dy, forward_dx), axis=-1)
+            flow_backward_raw = np.zeros_like(flow_forward_raw)
+            flow_forward = self.flow_transform(flow_forward_raw)
+            flow_backward = self.flow_transform(flow_backward_raw)
+
+            # update the mask for the pixels on the border
+            grid_x, grid_y = np.meshgrid(np.arange(self.image_size[0]), np.arange(self.image_size[1]), indexing="ij")
+            grid = np.stack((grid_y, grid_x), axis=-1)
+            grid_warp = grid + flow_forward_raw
+            location_y = grid_warp[:, :, 0].flatten()
+            location_x = grid_warp[:, :, 1].flatten()
+            I2_raw = np.array(I2_raw).astype(float)
+            I21_r = map_coordinates(I2_raw[:, :, 0], np.stack((location_x, location_y)), cval=-1).reshape(
+                (self.image_size[0], self.image_size[1])
+            )
+            I21_g = map_coordinates(I2_raw[:, :, 1], np.stack((location_x, location_y)), cval=-1).reshape(
+                (self.image_size[0], self.image_size[1])
+            )
+            I21_b = map_coordinates(I2_raw[:, :, 2], np.stack((location_x, location_y)), cval=-1).reshape(
+                (self.image_size[0], self.image_size[1])
+            )
+            I21_raw = np.stack((I21_r, I21_g, I21_b), axis=2)
+            mask = np.ones((self.image_size[0], self.image_size[1]))
+            mask[(I21_raw[:, :, 0] == -1) & (I21_raw[:, :, 1] == -1) & (I21_raw[:, :, 2] == -1)] = 0
+            mask[abs(I21_raw - I1_raw).sum(axis=-1) > 50] = 0
+            mask = self.ToTensor(mask)
+
+            if np.random.random() < self.real_reference_probability:
+                I_reference_output = I_reference_video_real
+                placeholder = torch.zeros_like(I1)
+                self_ref_flag = torch.zeros_like(I1)
+            else:
+                I_reference_output = I_reference_video
+                placeholder = I2 if np.random.random() < self.nonzero_placeholder_probability else torch.zeros_like(I1)
+                self_ref_flag = torch.ones_like(I1)
+
+            outputs = [
+                I1,
+                I2,
+                I_reference_output,
+                flow_forward,
+                flow_backward,
+                mask,
+                placeholder,
+                self_ref_flag,
+            ]
+
+        except Exception as e:
+            if combo_path is not None:
+                print("problem in ", combo_path)
+            print("problem in, ", image_a_path)
+            print(e)
+            return self.__getitem__(np.random.randint(0, len(self.image_pairs)))
+        return outputs
+
+    def __len__(self):
+        return len(self.image_pairs)

models/ContextualLoss.py

@@ -0,0 +1,303 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from utils.util import feature_normalize
+
+postpa = torchvision.transforms.Compose(
+    [
+        torchvision.transforms.Lambda(lambda x: x.mul_(1.0 / 255)),
+        torchvision.transforms.Normalize(
+            mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1]  # add imagenet mean
+        ),
+        torchvision.transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]),  # turn to RGB
+    ]
+)
+postpb = torchvision.transforms.Compose([torchvision.transforms.ToPILImage()])
+
+
+def post_processing(tensor):
+    t = postpa(tensor)  # denormalize the image since the optimized tensor is the normalized one
+    t[t > 1] = 1
+    t[t < 0] = 0
+    img = postpb(t)
+    img = np.array(img)
+    return img
+
+
+class ContextualLoss(nn.Module):
+    """
+    input is Al, Bl, channel = 1, range ~ [0, 255]
+    """
+
+    def __init__(self):
+        super(ContextualLoss, self).__init__()
+        return None
+
+    def forward(self, X_features, Y_features, h=0.1, feature_centering=True):
+        """
+        X_features&Y_features are are feature vectors or feature 2d array
+        h: bandwidth
+        return the per-sample loss
+        """
+        batch_size = X_features.shape[0]
+        feature_depth = X_features.shape[1]
+        feature_size = X_features.shape[2]
+
+        # to normalized feature vectors
+        if feature_centering:
+            X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
+                dim=-1
+            ).unsqueeze(dim=-1)
+            Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
+                dim=-1
+            ).unsqueeze(dim=-1)
+        X_features = feature_normalize(X_features).view(
+            batch_size, feature_depth, -1
+        )  # batch_size * feature_depth * feature_size^2
+        Y_features = feature_normalize(Y_features).view(
+            batch_size, feature_depth, -1
+        )  # batch_size * feature_depth * feature_size^2
+
+        # conine distance = 1 - similarity
+        X_features_permute = X_features.permute(0, 2, 1)  # batch_size * feature_size^2 * feature_depth
+        d = 1 - torch.matmul(X_features_permute, Y_features)  # batch_size * feature_size^2 * feature_size^2
+
+        # normalized distance: dij_bar
+        d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5)  # batch_size * feature_size^2 * feature_size^2
+
+        # pairwise affinity
+        w = torch.exp((1 - d_norm) / h)
+        A_ij = w / torch.sum(w, dim=-1, keepdim=True)
+
+        # contextual loss per sample
+        CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1)
+        return -torch.log(CX)
+
+
+class ContextualLoss_forward(nn.Module):
+    """
+    input is Al, Bl, channel = 1, range ~ [0, 255]
+    """
+
+    def __init__(self):
+        super(ContextualLoss_forward, self).__init__()
+        return None
+
+    def forward(self, X_features, Y_features, h=0.1, feature_centering=True):
+        """
+        X_features&Y_features are are feature vectors or feature 2d array
+        h: bandwidth
+        return the per-sample loss
+        """
+        batch_size = X_features.shape[0]
+        feature_depth = X_features.shape[1]
+        feature_size = X_features.shape[2]
+
+        # to normalized feature vectors
+        if feature_centering:
+            X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
+                dim=-1
+            ).unsqueeze(dim=-1)
+            Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
+                dim=-1
+            ).unsqueeze(dim=-1)
+        X_features = feature_normalize(X_features).view(
+            batch_size, feature_depth, -1
+        )  # batch_size * feature_depth * feature_size^2
+        Y_features = feature_normalize(Y_features).view(
+            batch_size, feature_depth, -1
+        )  # batch_size * feature_depth * feature_size^2
+
+        # conine distance = 1 - similarity
+        X_features_permute = X_features.permute(0, 2, 1)  # batch_size * feature_size^2 * feature_depth
+        d = 1 - torch.matmul(X_features_permute, Y_features)  # batch_size * feature_size^2 * feature_size^2
+
+        # normalized distance: dij_bar
+        d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5)  # batch_size * feature_size^2 * feature_size^2
+
+        # pairwise affinity
+        w = torch.exp((1 - d_norm) / h)
+        A_ij = w / torch.sum(w, dim=-1, keepdim=True)
+
+        # contextual loss per sample
+        CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1)
+        return -torch.log(CX)
+
+
+class ContextualLoss_complex(nn.Module):
+    """
+    input is Al, Bl, channel = 1, range ~ [0, 255]
+    """
+
+    def __init__(self):
+        super(ContextualLoss_complex, self).__init__()
+        return None
+
+    def forward(self, X_features, Y_features, h=0.1, patch_size=1, direction="forward"):
+        """
+        X_features&Y_features are are feature vectors or feature 2d array
+        h: bandwidth
+        return the per-sample loss
+        """
+        batch_size = X_features.shape[0]
+        feature_depth = X_features.shape[1]
+        feature_size = X_features.shape[2]
+
+        # to normalized feature vectors
+        X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
+            dim=-1
+        ).unsqueeze(dim=-1)
+        Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
+            dim=-1
+        ).unsqueeze(dim=-1)
+        X_features = feature_normalize(X_features)  # batch_size * feature_depth * feature_size^2
+        Y_features = feature_normalize(Y_features)  # batch_size * feature_depth * feature_size^2
+
+        # to normalized feature vectors
+        X_features = F.unfold(
+            X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
+        )  # batch_size * feature_depth_new * feature_size^2
+        Y_features = F.unfold(
+            Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
+        )  # batch_size * feature_depth_new * feature_size^2
+
+        # conine distance = 1 - similarity
+        X_features_permute = X_features.permute(0, 2, 1)  # batch_size * feature_size^2 * feature_depth
+        d = 1 - torch.matmul(X_features_permute, Y_features)  # batch_size * feature_size^2 * feature_size^2
+
+        # normalized distance: dij_bar
+        d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5)  # batch_size * feature_size^2 * feature_size^2
+
+        # pairwise affinity
+        w = torch.exp((1 - d_norm) / h)
+        A_ij = w / torch.sum(w, dim=-1, keepdim=True)
+
+        # contextual loss per sample
+        if direction == "forward":
+            CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1)
+        else:
+            CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1)
+
+        return -torch.log(CX)
+
+
+class ChamferDistance_patch_loss(nn.Module):
+    """
+    input is Al, Bl, channel = 1, range ~ [0, 255]
+    """
+
+    def __init__(self):
+        super(ChamferDistance_patch_loss, self).__init__()
+        return None
+
+    def forward(self, X_features, Y_features, patch_size=3, image_x=None, image_y=None, h=0.1, Y_features_in=None):
+        """
+        X_features&Y_features are are feature vectors or feature 2d array
+        h: bandwidth
+        return the per-sample loss
+        """
+        batch_size = X_features.shape[0]
+        feature_depth = X_features.shape[1]
+        feature_size = X_features.shape[2]
+
+        # to normalized feature vectors
+        X_features = F.unfold(
+            X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
+        )  # batch_size, feature_depth_new * feature_size^2
+        Y_features = F.unfold(
+            Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
+        )  # batch_size, feature_depth_new * feature_size^2
+
+        if image_x is not None and image_y is not None:
+            image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view(
+                batch_size, 3, -1
+            )
+            image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view(
+                batch_size, 3, -1
+            )
+
+        X_features_permute = X_features.permute(0, 2, 1)  # batch_size * feature_size^2 * feature_depth
+        similarity_matrix = torch.matmul(X_features_permute, Y_features)  # batch_size * feature_size^2 * feature_size^2
+        NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze()
+
+        if Y_features_in is not None:
+            loss = torch.mean((X_features - Y_features_in.detach()) ** 2)
+            Y_features_in = Y_features_in.detach()
+        else:
+            loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2)
+            Y_features_in = Y_features[:, :, NN_index].detach()
+
+        # re-arrange image
+        if image_x is not None and image_y is not None:
+            image_y_rearrange = image_y[:, :, NN_index]
+            image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size)
+            image_x = image_x.view(batch_size, 3, feature_size, feature_size)
+            image_y = image_y.view(batch_size, 3, feature_size, feature_size)
+
+        return loss
+
+
+class ChamferDistance_loss(nn.Module):
+    """
+    input is Al, Bl, channel = 1, range ~ [0, 255]
+    """
+
+    def __init__(self):
+        super(ChamferDistance_loss, self).__init__()
+        return None
+
+    def forward(self, X_features, Y_features, image_x, image_y, h=0.1, Y_features_in=None):
+        """
+        X_features&Y_features are are feature vectors or feature 2d array
+        h: bandwidth
+        return the per-sample loss
+        """
+        batch_size = X_features.shape[0]
+        feature_depth = X_features.shape[1]
+        feature_size = X_features.shape[2]
+
+        # to normalized feature vectors
+        X_features = feature_normalize(X_features).view(
+            batch_size, feature_depth, -1
+        )  # batch_size * feature_depth * feature_size^2
+        Y_features = feature_normalize(Y_features).view(
+            batch_size, feature_depth, -1
+        )  # batch_size * feature_depth * feature_size^2
+        image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view(
+            batch_size, 3, -1
+        )
+        image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view(
+            batch_size, 3, -1
+        )
+
+        X_features_permute = X_features.permute(0, 2, 1)  # batch_size * feature_size^2 * feature_depth
+        similarity_matrix = torch.matmul(X_features_permute, Y_features)  # batch_size * feature_size^2 * feature_size^2
+        NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze()
+        if Y_features_in is not None:
+            loss = torch.mean((X_features - Y_features_in.detach()) ** 2)
+            Y_features_in = Y_features_in.detach()
+        else:
+            loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2)
+            Y_features_in = Y_features[:, :, NN_index].detach()
+
+        # re-arrange image
+        image_y_rearrange = image_y[:, :, NN_index]
+        image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size)
+        image_x = image_x.view(batch_size, 3, feature_size, feature_size)
+        image_y = image_y.view(batch_size, 3, feature_size, feature_size)
+
+        return loss, Y_features_in, X_features
+
+
+if __name__ == "__main__":
+    contextual_loss = ContextualLoss()
+    batch_size = 32
+    feature_depth = 8
+    feature_size = 16
+    X_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size)
+    Y_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size)
+
+    cx_loss = contextual_loss(X_features, Y_features, 1)
+    print(cx_loss)

models/GAN_models.py

@@ -101,61 +101,13 @@ class Self_Attention(nn.Module):
         return out
 
 
-class Discriminator_WITHOUT_FC_x64_video(nn.Module):
+class Discriminator_x64(nn.Module):
     """
     Discriminative Network
     """
 
     def __init__(self, in_size=6, ndf=64):
-        super(Discriminator_WITHOUT_FC_x64_video, self).__init__()
-        self.in_size = in_size
-        self.ndf = ndf
-
-        self.layer1 = nn.Sequential(
-            # input size is in_size x 64 x 64
-            SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
-            nn.LeakyReLU(0.2, inplace=True),
-        )
-        self.layer2 = nn.Sequential(
-            # state size: ndf x 32 x 32
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 2),
-            nn.LeakyReLU(0.2, inplace=True),
-        )
-        self.layer3 = nn.Sequential(
-            # state size: (ndf * 2) x 16 x 16
-            SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 4),
-            nn.LeakyReLU(0.2, inplace=True),
-        )
-        self.layer4 = nn.Sequential(
-            # state size: (ndf * 4) x 8 x 8
-            # Self_Attention(self.ndf * 4, 'relu'),
-            SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 8),
-            nn.LeakyReLU(0.2, inplace=True),
-        )
-
-        self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0))
-
-    def forward(self, input):
-        feature1 = self.layer1(input)
-        feature2 = self.layer2(feature1)
-        feature3 = self.layer3(feature2)
-        feature4 = self.layer4(feature3)
-        output = self.last(feature4)
-        output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
-
-        return output, feature4
-
-
-class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module):
-    """
-    Discriminative Network
-    """
-
-    def __init__(self, in_size=6, ndf=64):
-        super(Discriminator_WITHOUT_FC_x64_video_BIG, self).__init__()
+        super(Discriminator_x64, self).__init__()
         self.in_size = in_size
         self.ndf = ndf
 
@@ -205,189 +157,237 @@ class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module):
         return output, feature4
 
 
-class Discriminator_WITHOUT_FC_x64(nn.Module):
-    """
-    Discriminative Network
-    """
+# class Discriminator_WITHOUT_FC_x64_video(nn.Module):
+#     """
+#     Discriminative Network
+#     """
 
-    def __init__(self, in_size=3, ndf=64):
-        super(Discriminator_WITHOUT_FC_x64, self).__init__()
-        self.in_size = in_size
-        self.ndf = ndf
+#     def __init__(self, in_size=6, ndf=64):
+#         super(Discriminator_WITHOUT_FC_x64_video, self).__init__()
+#         self.in_size = in_size
+#         self.ndf = ndf
 
-        self.main = nn.Sequential(
-            # input size is in_size x 64 x 64
-            SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: ndf x 32 x 32
-            Self_Attention(self.ndf),
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 2),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: (ndf * 2) x 16 x 16
-            SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 4),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: (ndf * 4) x 8 x 8
-            SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 8),
-            nn.LeakyReLU(0.2, inplace=True),
-            # sate size: (ndf * 8) x 4 x 4
-        )
+#         self.layer1 = nn.Sequential(
+#             # input size is in_size x 64 x 64
+#             SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
+#             nn.LeakyReLU(0.2, inplace=True),
+#         )
+#         self.layer2 = nn.Sequential(
+#             # state size: ndf x 32 x 32
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 2),
+#             nn.LeakyReLU(0.2, inplace=True),
+#         )
+#         self.layer3 = nn.Sequential(
+#             # state size: (ndf * 2) x 16 x 16
+#             SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 4),
+#             nn.LeakyReLU(0.2, inplace=True),
+#         )
+#         self.layer4 = nn.Sequential(
+#             # state size: (ndf * 4) x 8 x 8
+#             # Self_Attention(self.ndf * 4, 'relu'),
+#             SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 8),
+#             nn.LeakyReLU(0.2, inplace=True),
+#         )
 
-        self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
+#         self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0))
 
-    def forward(self, input):
+#     def forward(self, input):
+#         feature1 = self.layer1(input)
+#         feature2 = self.layer2(feature1)
+#         feature3 = self.layer3(feature2)
+#         feature4 = self.layer4(feature3)
+#         output = self.last(feature4)
+#         output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
 
-        feature = self.main(input)
-        output = self.last(feature)
-        output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
-
-        return output, feature
+#         return output, feature4
 
 
-class Discriminator_WITHOUT_FC_x64_BIG(nn.Module):
-    """
-    Discriminative Network
-    """
+# class Discriminator_WITHOUT_FC_x64(nn.Module):
+#     """
+#     Discriminative Network
+#     """
 
-    def __init__(self, in_size=3, ndf=64):
-        super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__()
-        self.in_size = in_size
-        self.ndf = ndf
+#     def __init__(self, in_size=3, ndf=64):
+#         super(Discriminator_WITHOUT_FC_x64, self).__init__()
+#         self.in_size = in_size
+#         self.ndf = ndf
 
-        self.main = nn.Sequential(
-            # 256 X 256
-            SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
-            nn.LeakyReLU(0.2, inplace=True),
-            # 128 X 128
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf),
-            nn.LeakyReLU(0.2, inplace=True),
-            # 64 X 64
-            Self_Attention(self.ndf),
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 2),
-            nn.LeakyReLU(0.2, inplace=True),
-            # 32 X 32
-            SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 4),
-            nn.LeakyReLU(0.2, inplace=True),
-            # 16 X 16
-            SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 8),
-            nn.LeakyReLU(0.2, inplace=True),
-            # 8 X 8
-            SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 16),
-            nn.LeakyReLU(0.2, inplace=True),
-        )
-        # 4 X 4
-        self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0))
+#         self.main = nn.Sequential(
+#             # input size is in_size x 64 x 64
+#             SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: ndf x 32 x 32
+#             Self_Attention(self.ndf),
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 2),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: (ndf * 2) x 16 x 16
+#             SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 4),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: (ndf * 4) x 8 x 8
+#             SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 8),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # sate size: (ndf * 8) x 4 x 4
+#         )
 
-    def forward(self, input):
+#         self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
 
-        feature = self.main(input)
-        output = self.last(feature)
-        output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
+#     def forward(self, input):
 
-        return output, feature
+#         feature = self.main(input)
+#         output = self.last(feature)
+#         output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
+
+#         return output, feature
 
 
-class Discriminator_WITHOUT_FC_x128(nn.Module):
-    """
-    Discriminative Network
-    """
+# class Discriminator_WITHOUT_FC_x64_BIG(nn.Module):
+#     """
+#     Discriminative Network
+#     """
 
-    def __init__(self, in_size=3, ndf=64):
-        super(Discriminator_WITHOUT_FC_x128, self).__init__()
-        self.in_size = in_size
-        self.ndf = ndf
+#     def __init__(self, in_size=3, ndf=64):
+#         super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__()
+#         self.in_size = in_size
+#         self.ndf = ndf
 
-        self.main = nn.Sequential(
-            # input size is in_size x 128 x 128
-            SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: ndf x 64 x 64
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: ndf x 32 x 32
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 2),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: (ndf * 24) x 16 x 16
-            SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 4),
-            nn.LeakyReLU(0.2, inplace=True),
-            # sate size: (ndf * 4) x 8 x 8
-            SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 8),
-            nn.LeakyReLU(0.2, inplace=True),
-            # sate size: (ndf * 8) x 4 x 4
-        )
+#         self.main = nn.Sequential(
+#             # 256 X 256
+#             SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # 128 X 128
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # 64 X 64
+#             Self_Attention(self.ndf),
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 2),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # 32 X 32
+#             SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 4),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # 16 X 16
+#             SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 8),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # 8 X 8
+#             SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 16),
+#             nn.LeakyReLU(0.2, inplace=True),
+#         )
+#         # 4 X 4
+#         self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0))
 
-        self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
+#     def forward(self, input):
 
-    def forward(self, input):
+#         feature = self.main(input)
+#         output = self.last(feature)
+#         output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
 
-        feature = self.main(input)
-        output = self.last(feature)
-
-        return output, feature
+#         return output, feature
 
 
-class Discriminator_WITHOUT_FC_x256(nn.Module):
-    """
-    Discriminative Network
-    """
+# class Discriminator_WITHOUT_FC_x128(nn.Module):
+#     """
+#     Discriminative Network
+#     """
 
-    def __init__(self, in_size=3, ndf=64):
-        super(Discriminator_WITHOUT_FC_x256, self).__init__()
-        self.in_size = in_size
-        self.ndf = ndf
+#     def __init__(self, in_size=3, ndf=64):
+#         super(Discriminator_WITHOUT_FC_x128, self).__init__()
+#         self.in_size = in_size
+#         self.ndf = ndf
 
-        self.main = nn.Sequential(
-            # input size is in_size x 256 x 256
-            SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: ndf x 128 x 128
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: ndf x 64 x 64
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: (ndf * 1) x 32 x 32
-            SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 2),
-            nn.LeakyReLU(0.2, inplace=True),
-            # state size: (ndf * 2) x 16 x 16
-            SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 4),
-            nn.LeakyReLU(0.2, inplace=True),
-            # sate size: (ndf * 4) x 8 x 8
-            SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
-            nn.InstanceNorm2d(self.ndf * 8),
-            nn.LeakyReLU(0.2, inplace=True),
-            # sate size: (ndf * 8) x 4 x 4
-        )
+#         self.main = nn.Sequential(
+#             # input size is in_size x 128 x 128
+#             SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: ndf x 64 x 64
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: ndf x 32 x 32
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 2),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: (ndf * 24) x 16 x 16
+#             SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 4),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # sate size: (ndf * 4) x 8 x 8
+#             SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 8),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # sate size: (ndf * 8) x 4 x 4
+#         )
 
-        self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
+#         self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
 
-    def forward(self, input):
+#     def forward(self, input):
 
-        feature = self.main(input)
-        output = self.last(feature)
+#         feature = self.main(input)
+#         output = self.last(feature)
 
-        return output, feature
+#         return output, feature
 
 
-if __name__ == "__main__":
-    discriminator1 = Discriminator_WITHOUT_FC_x64()
-    discriminator2 = Discriminator_WITHOUT_FC_x128()
-    discriminator3 = Discriminator_WITHOUT_FC_x256()
-    print(discriminator1)
-    print(discriminator2)
-    print(discriminator3)
+# class Discriminator_WITHOUT_FC_x256(nn.Module):
+#     """
+#     Discriminative Network
+#     """
+
+#     def __init__(self, in_size=3, ndf=64):
+#         super(Discriminator_WITHOUT_FC_x256, self).__init__()
+#         self.in_size = in_size
+#         self.ndf = ndf
+
+#         self.main = nn.Sequential(
+#             # input size is in_size x 256 x 256
+#             SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: ndf x 128 x 128
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: ndf x 64 x 64
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: (ndf * 1) x 32 x 32
+#             SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 2),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # state size: (ndf * 2) x 16 x 16
+#             SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 4),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # sate size: (ndf * 4) x 8 x 8
+#             SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
+#             nn.InstanceNorm2d(self.ndf * 8),
+#             nn.LeakyReLU(0.2, inplace=True),
+#             # sate size: (ndf * 8) x 4 x 4
+#         )
+
+#         self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
+
+#     def forward(self, input):
+
+#         feature = self.main(input)
+#         output = self.last(feature)
+
+#         return output, feature
+
+
+# if __name__ == "__main__":
+#     discriminator1 = Discriminator_WITHOUT_FC_x64()
+#     discriminator2 = Discriminator_WITHOUT_FC_x128()
+#     discriminator3 = Discriminator_WITHOUT_FC_x256()
+#     print(discriminator1)
+#     print(discriminator2)
+#     print(discriminator3)

models/NonlocalNet.py

@@ -189,13 +189,13 @@ class Self_Attn(nn.Module):
         return out
 
 
-class VGG19_feature_color_torchversion(nn.Module):
+class VGG19_pytorch(nn.Module):
     """
     NOTE: no need to pre-process the input; input tensor should range in [0,1]
     """
 
     def __init__(self, pool="max"):
-        super(VGG19_feature_color_torchversion, self).__init__()
+        super(VGG19_pytorch, self).__init__()
         self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
         self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
         self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)

opt.txt

@@ -0,0 +1,37 @@
+train.py 2020-12-21 15:03:17
+------------ Options -------------
+batch_size: 16
+beta1: 0.5
+checkpoint_dir: checkpoints/default
+contextual_loss_direction: forward
+data_root: /home/zhanbo/remote/video/video_pair3/
+data_root_imagenet: /home/zhanbo/remote/imagenet_ref_pair/
+domain_invariant: False
+epoch: 40
+gpu_ids: 0,1,2,3
+ic: 7
+image_size: [216, 384]
+load_pretrained_model: True
+lr: 1e-05
+lr_gamma: 0.1
+lr_step: 100
+luminance_noise: 2.0
+nonzero_placeholder_probability: 0.0
+permute_data: True
+print_step: 2
+real_reference_probability: 0.7
+resume: False
+resume_epoch: 0
+tb_log_step: 50
+weight_consistent: 0.02
+weight_contextual: 0.2
+weight_gan: 0.2
+weight_nonlocal_consistent: 0.0
+weight_nonlocal_smoothness: 0.0
+weight_perceptual: 0.001
+weight_smoothness: 5.0
+weigth_l1: 2.0
+with_bad: True
+with_mid: True
+workers: 8
+-------------- End ----------------

requirements.txt

@@ -14,4 +14,5 @@ prefetch_generator
 PyYAML
 tensorboard
 future
-numba
+numba
+pypng

tensorboardX/__init__.py

@@ -0,0 +1,5 @@
+"""A module for visualization with tensorboard
+"""
+
+from .writer import FileWriter, SummaryWriter
+from .record_writer import RecordWriter

tensorboardX/crc32c.py

@@ -0,0 +1,123 @@
+import array
+
+
+CRC_TABLE = (
+    0x00000000, 0xf26b8303, 0xe13b70f7, 0x1350f3f4,
+    0xc79a971f, 0x35f1141c, 0x26a1e7e8, 0xd4ca64eb,
+    0x8ad958cf, 0x78b2dbcc, 0x6be22838, 0x9989ab3b,
+    0x4d43cfd0, 0xbf284cd3, 0xac78bf27, 0x5e133c24,
+    0x105ec76f, 0xe235446c, 0xf165b798, 0x030e349b,
+    0xd7c45070, 0x25afd373, 0x36ff2087, 0xc494a384,
+    0x9a879fa0, 0x68ec1ca3, 0x7bbcef57, 0x89d76c54,
+    0x5d1d08bf, 0xaf768bbc, 0xbc267848, 0x4e4dfb4b,
+    0x20bd8ede, 0xd2d60ddd, 0xc186fe29, 0x33ed7d2a,
+    0xe72719c1, 0x154c9ac2, 0x061c6936, 0xf477ea35,
+    0xaa64d611, 0x580f5512, 0x4b5fa6e6, 0xb93425e5,
+    0x6dfe410e, 0x9f95c20d, 0x8cc531f9, 0x7eaeb2fa,
+    0x30e349b1, 0xc288cab2, 0xd1d83946, 0x23b3ba45,
+    0xf779deae, 0x05125dad, 0x1642ae59, 0xe4292d5a,
+    0xba3a117e, 0x4851927d, 0x5b016189, 0xa96ae28a,
+    0x7da08661, 0x8fcb0562, 0x9c9bf696, 0x6ef07595,
+    0x417b1dbc, 0xb3109ebf, 0xa0406d4b, 0x522bee48,
+    0x86e18aa3, 0x748a09a0, 0x67dafa54, 0x95b17957,
+    0xcba24573, 0x39c9c670, 0x2a993584, 0xd8f2b687,
+    0x0c38d26c, 0xfe53516f, 0xed03a29b, 0x1f682198,
+    0x5125dad3, 0xa34e59d0, 0xb01eaa24, 0x42752927,
+    0x96bf4dcc, 0x64d4cecf, 0x77843d3b, 0x85efbe38,
+    0xdbfc821c, 0x2997011f, 0x3ac7f2eb, 0xc8ac71e8,
+    0x1c661503, 0xee0d9600, 0xfd5d65f4, 0x0f36e6f7,
+    0x61c69362, 0x93ad1061, 0x80fde395, 0x72966096,
+    0xa65c047d, 0x5437877e, 0x4767748a, 0xb50cf789,
+    0xeb1fcbad, 0x197448ae, 0x0a24bb5a, 0xf84f3859,
+    0x2c855cb2, 0xdeeedfb1, 0xcdbe2c45, 0x3fd5af46,
+    0x7198540d, 0x83f3d70e, 0x90a324fa, 0x62c8a7f9,
+    0xb602c312, 0x44694011, 0x5739b3e5, 0xa55230e6,
+    0xfb410cc2, 0x092a8fc1, 0x1a7a7c35, 0xe811ff36,
+    0x3cdb9bdd, 0xceb018de, 0xdde0eb2a, 0x2f8b6829,
+    0x82f63b78, 0x709db87b, 0x63cd4b8f, 0x91a6c88c,
+    0x456cac67, 0xb7072f64, 0xa457dc90, 0x563c5f93,
+    0x082f63b7, 0xfa44e0b4, 0xe9141340, 0x1b7f9043,
+    0xcfb5f4a8, 0x3dde77ab, 0x2e8e845f, 0xdce5075c,
+    0x92a8fc17, 0x60c37f14, 0x73938ce0, 0x81f80fe3,
+    0x55326b08, 0xa759e80b, 0xb4091bff, 0x466298fc,
+    0x1871a4d8, 0xea1a27db, 0xf94ad42f, 0x0b21572c,
+    0xdfeb33c7, 0x2d80b0c4, 0x3ed04330, 0xccbbc033,
+    0xa24bb5a6, 0x502036a5, 0x4370c551, 0xb11b4652,
+    0x65d122b9, 0x97baa1ba, 0x84ea524e, 0x7681d14d,
+    0x2892ed69, 0xdaf96e6a, 0xc9a99d9e, 0x3bc21e9d,
+    0xef087a76, 0x1d63f975, 0x0e330a81, 0xfc588982,
+    0xb21572c9, 0x407ef1ca, 0x532e023e, 0xa145813d,
+    0x758fe5d6, 0x87e466d5, 0x94b49521, 0x66df1622,
+    0x38cc2a06, 0xcaa7a905, 0xd9f75af1, 0x2b9cd9f2,
+    0xff56bd19, 0x0d3d3e1a, 0x1e6dcdee, 0xec064eed,
+    0xc38d26c4, 0x31e6a5c7, 0x22b65633, 0xd0ddd530,
+    0x0417b1db, 0xf67c32d8, 0xe52cc12c, 0x1747422f,
+    0x49547e0b, 0xbb3ffd08, 0xa86f0efc, 0x5a048dff,
+    0x8ecee914, 0x7ca56a17, 0x6ff599e3, 0x9d9e1ae0,
+    0xd3d3e1ab, 0x21b862a8, 0x32e8915c, 0xc083125f,
+    0x144976b4, 0xe622f5b7, 0xf5720643, 0x07198540,
+    0x590ab964, 0xab613a67, 0xb831c993, 0x4a5a4a90,
+    0x9e902e7b, 0x6cfbad78, 0x7fab5e8c, 0x8dc0dd8f,
+    0xe330a81a, 0x115b2b19, 0x020bd8ed, 0xf0605bee,
+    0x24aa3f05, 0xd6c1bc06, 0xc5914ff2, 0x37faccf1,
+    0x69e9f0d5, 0x9b8273d6, 0x88d28022, 0x7ab90321,
+    0xae7367ca, 0x5c18e4c9, 0x4f48173d, 0xbd23943e,
+    0xf36e6f75, 0x0105ec76, 0x12551f82, 0xe03e9c81,
+    0x34f4f86a, 0xc69f7b69, 0xd5cf889d, 0x27a40b9e,
+    0x79b737ba, 0x8bdcb4b9, 0x988c474d, 0x6ae7c44e,
+    0xbe2da0a5, 0x4c4623a6, 0x5f16d052, 0xad7d5351,
+)
+
+
+CRC_INIT = 0
+
+_MASK = 0xFFFFFFFF
+
+
+def crc_update(crc, data):
+    """Update CRC-32C checksum with data.
+
+    Args:
+      crc: 32-bit checksum to update as long.
+      data: byte array, string or iterable over bytes.
+
+    Returns:
+      32-bit updated CRC-32C as long.
+    """
+
+    if type(data) != array.array or data.itemsize != 1:
+        buf = array.array("B", data)
+    else:
+        buf = data
+
+    crc ^= _MASK
+    for b in buf:
+        table_index = (crc ^ b) & 0xff
+        crc = (CRC_TABLE[table_index] ^ (crc >> 8)) & _MASK
+    return crc ^ _MASK
+
+
+def crc_finalize(crc):
+    """Finalize CRC-32C checksum.
+
+    This function should be called as last step of crc calculation.
+
+    Args:
+      crc: 32-bit checksum as long.
+
+    Returns:
+      finalized 32-bit checksum as long
+    """
+    return crc & _MASK
+
+
+def crc32c(data):
+    """Compute CRC-32C checksum of the data.
+
+    Args:
+      data: byte array, string or iterable over bytes.
+
+    Returns:
+      32-bit CRC-32C checksum of data as long.
+    """
+    return crc_finalize(crc_update(CRC_INIT, data))

tensorboardX/embedding.py

@@ -0,0 +1,49 @@
+import os
+
+
+def make_tsv(metadata, save_path):
+    metadata = [str(x) for x in metadata]
+    with open(os.path.join(save_path, 'metadata.tsv'), 'w') as f:
+        for x in metadata:
+            f.write(x + '\n')
+
+
+# https://github.com/tensorflow/tensorboard/issues/44 image label will be squared
+def make_sprite(label_img, save_path):
+    import math
+    import torch
+    import torchvision
+    from .x2num import makenp
+    # this ensures the sprite image has correct dimension as described in
+    # https://www.tensorflow.org/get_started/embedding_viz
+    nrow = int(math.ceil((label_img.size(0)) ** 0.5))
+
+    label_img = torch.from_numpy(makenp(label_img))  # for other framework
+    # augment images so that #images equals nrow*nrow
+    label_img = torch.cat((label_img, torch.randn(nrow ** 2 - label_img.size(0), *label_img.size()[1:]) * 255), 0)
+
+    torchvision.utils.save_image(label_img, os.path.join(save_path, 'sprite.png'), nrow=nrow, padding=0)
+
+
+def append_pbtxt(metadata, label_img, save_path, global_step, tag):
+    with open(os.path.join(save_path, 'projector_config.pbtxt'), 'a') as f:
+        # step = os.path.split(save_path)[-1]
+        f.write('embeddings {\n')
+        f.write('tensor_name: "{}:{}"\n'.format(tag, global_step))
+        f.write('tensor_path: "{}"\n'.format(os.path.join(global_step, 'tensors.tsv')))
+        if metadata is not None:
+            f.write('metadata_path: "{}"\n'.format(os.path.join(global_step, 'metadata.tsv')))
+        if label_img is not None:
+            f.write('sprite {\n')
+            f.write('image_path: "{}"\n'.format(os.path.join(global_step, 'sprite.png')))
+            f.write('single_image_dim: {}\n'.format(label_img.size(3)))
+            f.write('single_image_dim: {}\n'.format(label_img.size(2)))
+            f.write('}\n')
+        f.write('}\n')
+
+
+def make_mat(matlist, save_path):
+    with open(os.path.join(save_path, 'tensors.tsv'), 'w') as f:
+        for x in matlist:
+            x = [str(i) for i in x]
+            f.write('\t'.join(x) + '\n')

tensorboardX/event_file_writer.py

@@ -0,0 +1,196 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Writes events to disk in a logdir."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import os.path
+import socket
+import threading
+import time
+
+import six
+
+from .src import event_pb2
+from .record_writer import RecordWriter
+
+
+def directory_check(path):
+    '''Initialize the directory for log files.'''
+    # If the direcotry does not exist, create it!
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+class EventsWriter(object):
+    '''Writes `Event` protocol buffers to an event file.'''
+
+    def __init__(self, file_prefix):
+        '''
+        Events files have a name of the form
+        '/some/file/path/events.out.tfevents.[timestamp].[hostname]'
+        '''
+        self._file_prefix = file_prefix + ".out.tfevents." + str(time.time())[:10] + "." + socket.gethostname()
+
+        # Open(Create) the log file with the particular form of name.
+        logging.basicConfig(filename=self._file_prefix)
+
+        self._num_outstanding_events = 0
+
+        self._py_recordio_writer = RecordWriter(self._file_prefix)
+
+        # Initialize an event instance.
+        self._event = event_pb2.Event()
+
+        self._event.wall_time = time.time()
+
+        self.write_event(self._event)
+
+    def write_event(self, event):
+        '''Append "event" to the file.'''
+
+        # Check if event is of type event_pb2.Event proto.
+        if not isinstance(event, event_pb2.Event):
+            raise TypeError("Expected an event_pb2.Event proto, "
+                            " but got %s" % type(event))
+        return self._write_serialized_event(event.SerializeToString())
+
+    def _write_serialized_event(self, event_str):
+        self._num_outstanding_events += 1
+        self._py_recordio_writer.write(event_str)
+
+    def flush(self):
+        '''Flushes the event file to disk.'''
+        self._num_outstanding_events = 0
+        return True
+
+    def close(self):
+        '''Call self.flush().'''
+        return_value = self.flush()
+        return return_value
+
+
+class EventFileWriter(object):
+    """Writes `Event` protocol buffers to an event file.
+    The `EventFileWriter` class creates an event file in the specified directory,
+    and asynchronously writes Event protocol buffers to the file. The Event file
+    is encoded using the tfrecord format, which is similar to RecordIO.
+    @@__init__
+    @@add_event
+    @@flush
+    @@close
+    """
+
+    def __init__(self, logdir, max_queue=10, flush_secs=120):
+        """Creates a `EventFileWriter` and an event file to write to.
+        On construction the summary writer creates a new event file in `logdir`.
+        This event file will contain `Event` protocol buffers, which are written to
+        disk via the add_event method.
+        The other arguments to the constructor control the asynchronous writes to
+        the event file:
+        *  `flush_secs`: How often, in seconds, to flush the added summaries
+           and events to disk.
+        *  `max_queue`: Maximum number of summaries or events pending to be
+           written to disk before one of the 'add' calls block.
+        Args:
+          logdir: A string. Directory where event file will be written.
+          max_queue: Integer. Size of the queue for pending events and summaries.
+          flush_secs: Number. How often, in seconds, to flush the
+            pending events and summaries to disk.
+        """
+        self._logdir = logdir
+        directory_check(self._logdir)
+        self._event_queue = six.moves.queue.Queue(max_queue)
+        self._ev_writer = EventsWriter(os.path.join(self._logdir, "events"))
+        self._closed = False
+        self._worker = _EventLoggerThread(self._event_queue, self._ev_writer,
+                                          flush_secs)
+
+        self._worker.start()
+
+    def get_logdir(self):
+        """Returns the directory where event file will be written."""
+        return self._logdir
+
+    def reopen(self):
+        """Reopens the EventFileWriter.
+        Can be called after `close()` to add more events in the same directory.
+        The events will go into a new events file.
+        Does nothing if the EventFileWriter was not closed.
+        """
+        if self._closed:
+            self._closed = False
+
+    def add_event(self, event):
+        """Adds an event to the event file.
+        Args:
+          event: An `Event` protocol buffer.
+        """
+        if not self._closed:
+            self._event_queue.put(event)
+
+    def flush(self):
+        """Flushes the event file to disk.
+        Call this method to make sure that all pending events have been written to
+        disk.
+        """
+        self._event_queue.join()
+        self._ev_writer.flush()
+
+    def close(self):
+        """Flushes the event file to disk and close the file.
+        Call this method when you do not need the summary writer anymore.
+        """
+        self.flush()
+        self._ev_writer.close()
+        self._closed = True
+
+
+class _EventLoggerThread(threading.Thread):
+    """Thread that logs events."""
+
+    def __init__(self, queue, ev_writer, flush_secs):
+        """Creates an _EventLoggerThread.
+        Args:
+          queue: A Queue from which to dequeue events.
+          ev_writer: An event writer. Used to log brain events for
+           the visualizer.
+          flush_secs: How often, in seconds, to flush the
+            pending file to disk.
+        """
+        threading.Thread.__init__(self)
+        self.daemon = True
+        self._queue = queue
+        self._ev_writer = ev_writer
+        self._flush_secs = flush_secs
+        # The first event will be flushed immediately.
+        self._next_event_flush_time = 0
+
+    def run(self):
+        while True:
+            event = self._queue.get()
+            try:
+                self._ev_writer.write_event(event)
+                # Flush the event writer every so often.
+                now = time.time()
+                if now > self._next_event_flush_time:
+                    self._ev_writer.flush()
+                    # Do it again in two minutes.
+                    self._next_event_flush_time = now + self._flush_secs
+            finally:
+                self._queue.task_done()

tensorboardX/graph.py

@@ -0,0 +1,66 @@
+from .src.graph_pb2 import GraphDef
+from .src.node_def_pb2 import NodeDef
+from .src.versions_pb2 import VersionDef
+from .src.attr_value_pb2 import AttrValue
+from .src.tensor_shape_pb2 import TensorShapeProto
+
+from distutils.version import LooseVersion
+
+
+def replace(name, scope):
+    return '/'.join([scope[name], name])
+
+
+def parse(graph):
+    import torch
+    scope = {}
+    for n in graph.nodes():
+        inputs = [i.uniqueName() for i in n.inputs()]
+        for i in range(1, len(inputs)):
+            scope[inputs[i]] = n.scopeName()
+
+        uname = next(iter(n.outputs())).uniqueName()
+        assert n.scopeName() != '', '{} has empty scope name'.format(n)
+        scope[uname] = n.scopeName()
+    if LooseVersion(torch.__version__) >= LooseVersion("0.4"):
+        scope['0'] = 'input'
+    else:
+        scope['1'] = 'input'
+
+    nodes = []
+    for n in graph.nodes():
+        attrs = {k: n[k] for k in n.attributeNames()}
+        attrs = str(attrs).replace("'", ' ')  # singlequote will be escaped by tensorboard
+        if any(i.uniqueName() not in scope.keys() for i in n.inputs()):  # 0.3.1 workaround
+            continue
+        inputs = [replace(i.uniqueName(), scope) for i in n.inputs()]
+        uname = next(iter(n.outputs())).uniqueName()  # FIXME: only first output is considered
+        nodes.append({'name': replace(uname, scope), 'op': n.kind(), 'inputs': inputs, 'attr': attrs})
+
+    for n in graph.inputs():
+        uname = n.uniqueName()
+        if uname not in scope.keys():
+            scope[uname] = 'unused'
+        nodes.append({'name': replace(uname, scope), 'op': 'Parameter', 'inputs': [], 'attr': str(n.type())})
+
+    return nodes
+
+
+def graph(model, args, verbose=False):
+    import torch
+    with torch.onnx.set_training(model, False):
+        trace, _ = torch.jit.trace(model, args)
+    if LooseVersion(torch.__version__) >= LooseVersion("0.4"):
+        torch.onnx._optimize_trace(trace, False)
+    else:
+        torch.onnx._optimize_trace(trace)
+    graph = trace.graph()
+    if verbose:
+        print(graph)
+    list_of_nodes = parse(graph)
+    nodes = []
+    for node in list_of_nodes:
+        nodes.append(
+            NodeDef(name=node['name'], op=node['op'], input=node['inputs'],
+                    attr={'lanpa': AttrValue(s=node['attr'].encode(encoding='utf_8'))}))
+    return GraphDef(node=nodes, versions=VersionDef(producer=22))

tensorboardX/graph_onnx.py

@@ -0,0 +1,158 @@
+from .src.graph_pb2 import GraphDef
+from .src.node_def_pb2 import NodeDef
+from .src.versions_pb2 import VersionDef
+from .src.attr_value_pb2 import AttrValue
+from .src.tensor_shape_pb2 import TensorShapeProto
+# from .src.onnx_pb2 import ModelProto
+
+
+def gg(fname):
+    import onnx  # 0.2.1
+    m = onnx.load(fname)
+    nodes_proto = []
+    nodes = []
+    g = m.graph
+    import itertools
+    for node in itertools.chain(g.input, g.output):
+        nodes_proto.append(node)
+
+    for node in nodes_proto:
+        shapeproto = TensorShapeProto(
+            dim=[TensorShapeProto.Dim(size=d.dim_value) for d in node.type.tensor_type.shape.dim])
+        nodes.append(NodeDef(
+            name=node.name,
+            op='Variable',
+            input=[],
+            attr={
+                'dtype': AttrValue(type=node.type.tensor_type.elem_type),
+                'shape': AttrValue(shape=shapeproto),
+            })
+        )
+
+    for node in g.node:
+        attr = []
+        for s in node.attribute:
+            attr.append(' = '.join([str(f[1]) for f in s.ListFields()]))
+        attr = ', '.join(attr).encode(encoding='utf_8')
+
+        nodes.append(NodeDef(
+            name=node.output[0],
+            op=node.op_type,
+            input=node.input,
+            attr={'parameters': AttrValue(s=attr)},
+        ))
+    # two pass token replacement, appends opname to object id
+    mapping = {}
+    for node in nodes:
+        mapping[node.name] = node.op + '_' + node.name
+
+    nodes, mapping = updatenodes(nodes, mapping)
+    mapping = smartGrouping(nodes, mapping)
+    nodes, mapping = updatenodes(nodes, mapping)
+
+    return GraphDef(node=nodes, versions=VersionDef(producer=22))
+
+
+def updatenodes(nodes, mapping):
+    for node in nodes:
+        newname = mapping[node.name]
+        node.name = newname
+        newinput = []
+        for inputnode in list(node.input):
+            newinput.append(mapping[inputnode])
+            node.input.remove(inputnode)
+        node.input.extend(newinput)
+    newmap = {}
+    for k, v in mapping.items():
+        newmap[v] = v
+    return nodes, newmap
+
+
+def findnode(nodes, name):
+    """ input: node name
+        returns: node object
+    """
+    for n in nodes:
+        if n.name == name:
+            return n
+
+
+def parser(s, nodes, node):
+    print(s)
+    if len(s) == 0:
+        return
+    if len(s) > 0:
+        if s[0] == node.op:
+            print(s[0], node.name, s[1], node.input)
+            for n in node.input:
+                print(n, s[1])
+                parser(s[1], nodes, findnode(nodes, n))
+        else:
+            return False
+
+
+# TODO: use recursive parse
+
+def smartGrouping(nodes, mapping):
+    # a Fully Conv is: (TODO: check var1.size(0)==var2.size(0))
+    # GEMM <-- Variable (c1)
+    #  ^-- Transpose (c2) <-- Variable (c3)
+
+    # a Conv with bias is: (TODO: check var1.size(0)==var2.size(0))
+    # Add <-- Conv (c2) <-- Variable (c3)
+    #  ^-- Variable (c1)
+    #
+    # gemm = ('Gemm', ('Variable', ('Transpose', ('Variable'))))
+
+    FCcounter = 1
+    Convcounter = 1
+    for node in nodes:
+        if node.op == 'Gemm':
+            c1 = c2 = c3 = False
+            for name_in in node.input:
+                n = findnode(nodes, name_in)
+                if n.op == 'Variable':
+                    c1 = True
+                    c1name = n.name
+                if n.op == 'Transpose':
+                    c2 = True
+                    c2name = n.name
+                    if len(n.input) == 1:
+                        nn = findnode(nodes, n.input[0])
+                        if nn.op == 'Variable':
+                            c3 = True
+                            c3name = nn.name
+                # print(n.op, n.name, c1, c2, c3)
+            if c1 and c2 and c3:
+                # print(c1name, c2name, c3name)
+                mapping[c1name] = 'FC{}/{}'.format(FCcounter, c1name)
+                mapping[c2name] = 'FC{}/{}'.format(FCcounter, c2name)
+                mapping[c3name] = 'FC{}/{}'.format(FCcounter, c3name)
+                mapping[node.name] = 'FC{}/{}'.format(FCcounter, node.name)
+                FCcounter += 1
+                continue
+        if node.op == 'Add':
+            c1 = c2 = c3 = False
+            for name_in in node.input:
+                n = findnode(nodes, name_in)
+                if n.op == 'Variable':
+                    c1 = True
+                    c1name = n.name
+                if n.op == 'Conv':
+                    c2 = True
+                    c2name = n.name
+                    if len(n.input) >= 1:
+                        for nn_name in n.input:
+                            nn = findnode(nodes, nn_name)
+                            if nn.op == 'Variable':
+                                c3 = True
+                                c3name = nn.name
+
+            if c1 and c2 and c3:
+                # print(c1name, c2name, c3name)
+                mapping[c1name] = 'Conv{}/{}'.format(Convcounter, c1name)
+                mapping[c2name] = 'Conv{}/{}'.format(Convcounter, c2name)
+                mapping[c3name] = 'Conv{}/{}'.format(Convcounter, c3name)
+                mapping[node.name] = 'Conv{}/{}'.format(Convcounter, node.name)
+                Convcounter += 1
+    return mapping

tensorboardX/record_writer.py

@@ -0,0 +1,47 @@
+"""
+To write tf_record into file. Here we use it for tensorboard's event writting.
+The code was borrow from https://github.com/TeamHG-Memex/tensorboard_logger
+"""
+
+import re
+import struct
+
+from .crc32c import crc32c
+
+_VALID_OP_NAME_START = re.compile('^[A-Za-z0-9.]')
+_VALID_OP_NAME_PART = re.compile('[A-Za-z0-9_.\\-/]+')
+
+
+class RecordWriter(object):
+    def __init__(self, path, flush_secs=2):
+        self._name_to_tf_name = {}
+        self._tf_names = set()
+        self.path = path
+        self.flush_secs = flush_secs  # TODO. flush every flush_secs, not every time.
+        self._writer = None
+        self._writer = open(path, 'wb')
+
+    def write(self, event_str):
+        w = self._writer.write
+        header = struct.pack('Q', len(event_str))
+        w(header)
+        w(struct.pack('I', masked_crc32c(header)))
+        w(event_str)
+        w(struct.pack('I', masked_crc32c(event_str)))
+        self._writer.flush()
+
+
+def masked_crc32c(data):
+    x = u32(crc32c(data))
+    return u32(((x >> 15) | u32(x << 17)) + 0xa282ead8)
+
+
+def u32(x):
+    return x & 0xffffffff
+
+
+def make_valid_tf_name(name):
+    if not _VALID_OP_NAME_START.match(name):
+        # Must make it valid somehow, but don't want to remove stuff
+        name = '.' + name
+    return '_'.join(_VALID_OP_NAME_PART.findall(name))

tensorboardX/src/attr_value_pb2.py

@@ -0,0 +1,372 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/attr_value.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorboardX.src import tensor_pb2 as tensorboardX_dot_src_dot_tensor__pb2
+from tensorboardX.src import tensor_shape_pb2 as tensorboardX_dot_src_dot_tensor__shape__pb2
+from tensorboardX.src import types_pb2 as tensorboardX_dot_src_dot_types__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/attr_value.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n!tensorboardX/src/attr_value.proto\x12\x0btensorboard\x1a\x1dtensorboardX/src/tensor.proto\x1a#tensorboardX/src/tensor_shape.proto\x1a\x1ctensorboardX/src/types.proto\"\xaf\x04\n\tAttrValue\x12\x0b\n\x01s\x18\x02 \x01(\x0cH\x00\x12\x0b\n\x01i\x18\x03 \x01(\x03H\x00\x12\x0b\n\x01\x66\x18\x04 \x01(\x02H\x00\x12\x0b\n\x01\x62\x18\x05 \x01(\x08H\x00\x12%\n\x04type\x18\x06 \x01(\x0e\x32\x15.tensorboard.DataTypeH\x00\x12.\n\x05shape\x18\x07 \x01(\x0b\x32\x1d.tensorboard.TensorShapeProtoH\x00\x12*\n\x06tensor\x18\x08 \x01(\x0b\x32\x18.tensorboard.TensorProtoH\x00\x12\x30\n\x04list\x18\x01 \x01(\x0b\x32 .tensorboard.AttrValue.ListValueH\x00\x12)\n\x04\x66unc\x18\n \x01(\x0b\x32\x19.tensorboard.NameAttrListH\x00\x12\x15\n\x0bplaceholder\x18\t \x01(\tH\x00\x1a\xed\x01\n\tListValue\x12\t\n\x01s\x18\x02 \x03(\x0c\x12\r\n\x01i\x18\x03 \x03(\x03\x42\x02\x10\x01\x12\r\n\x01\x66\x18\x04 \x03(\x02\x42\x02\x10\x01\x12\r\n\x01\x62\x18\x05 \x03(\x08\x42\x02\x10\x01\x12\'\n\x04type\x18\x06 \x03(\x0e\x32\x15.tensorboard.DataTypeB\x02\x10\x01\x12,\n\x05shape\x18\x07 \x03(\x0b\x32\x1d.tensorboard.TensorShapeProto\x12(\n\x06tensor\x18\x08 \x03(\x0b\x32\x18.tensorboard.TensorProto\x12\'\n\x04\x66unc\x18\t \x03(\x0b\x32\x19.tensorboard.NameAttrListB\x07\n\x05value\"\x94\x01\n\x0cNameAttrList\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04\x61ttr\x18\x02 \x03(\x0b\x32#.tensorboard.NameAttrList.AttrEntry\x1a\x43\n\tAttrEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.tensorboard.AttrValue:\x02\x38\x01\x42\x30\n\x18org.tensorflow.frameworkB\x0f\x41ttrValueProtosP\x01\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorboardX_dot_src_dot_tensor__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_tensor__shape__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_types__pb2.DESCRIPTOR,])
+
+
+
+
+_ATTRVALUE_LISTVALUE = _descriptor.Descriptor(
+  name='ListValue',
+  full_name='tensorboard.AttrValue.ListValue',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='s', full_name='tensorboard.AttrValue.ListValue.s', index=0,
+      number=2, type=12, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='i', full_name='tensorboard.AttrValue.ListValue.i', index=1,
+      number=3, type=3, cpp_type=2, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='f', full_name='tensorboard.AttrValue.ListValue.f', index=2,
+      number=4, type=2, cpp_type=6, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='b', full_name='tensorboard.AttrValue.ListValue.b', index=3,
+      number=5, type=8, cpp_type=7, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='type', full_name='tensorboard.AttrValue.ListValue.type', index=4,
+      number=6, type=14, cpp_type=8, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='shape', full_name='tensorboard.AttrValue.ListValue.shape', index=5,
+      number=7, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tensor', full_name='tensorboard.AttrValue.ListValue.tensor', index=6,
+      number=8, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='func', full_name='tensorboard.AttrValue.ListValue.func', index=7,
+      number=9, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=462,
+  serialized_end=699,
+)
+
+_ATTRVALUE = _descriptor.Descriptor(
+  name='AttrValue',
+  full_name='tensorboard.AttrValue',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='s', full_name='tensorboard.AttrValue.s', index=0,
+      number=2, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='i', full_name='tensorboard.AttrValue.i', index=1,
+      number=3, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='f', full_name='tensorboard.AttrValue.f', index=2,
+      number=4, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='b', full_name='tensorboard.AttrValue.b', index=3,
+      number=5, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='type', full_name='tensorboard.AttrValue.type', index=4,
+      number=6, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='shape', full_name='tensorboard.AttrValue.shape', index=5,
+      number=7, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tensor', full_name='tensorboard.AttrValue.tensor', index=6,
+      number=8, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='list', full_name='tensorboard.AttrValue.list', index=7,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='func', full_name='tensorboard.AttrValue.func', index=8,
+      number=10, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='placeholder', full_name='tensorboard.AttrValue.placeholder', index=9,
+      number=9, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[_ATTRVALUE_LISTVALUE, ],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='value', full_name='tensorboard.AttrValue.value',
+      index=0, containing_type=None, fields=[]),
+  ],
+  serialized_start=149,
+  serialized_end=708,
+)
+
+
+_NAMEATTRLIST_ATTRENTRY = _descriptor.Descriptor(
+  name='AttrEntry',
+  full_name='tensorboard.NameAttrList.AttrEntry',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='key', full_name='tensorboard.NameAttrList.AttrEntry.key', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='value', full_name='tensorboard.NameAttrList.AttrEntry.value', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=_descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001')),
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=792,
+  serialized_end=859,
+)
+
+_NAMEATTRLIST = _descriptor.Descriptor(
+  name='NameAttrList',
+  full_name='tensorboard.NameAttrList',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='name', full_name='tensorboard.NameAttrList.name', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='attr', full_name='tensorboard.NameAttrList.attr', index=1,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[_NAMEATTRLIST_ATTRENTRY, ],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=711,
+  serialized_end=859,
+)
+
+_ATTRVALUE_LISTVALUE.fields_by_name['type'].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE
+_ATTRVALUE_LISTVALUE.fields_by_name['shape'].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO
+_ATTRVALUE_LISTVALUE.fields_by_name['tensor'].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO
+_ATTRVALUE_LISTVALUE.fields_by_name['func'].message_type = _NAMEATTRLIST
+_ATTRVALUE_LISTVALUE.containing_type = _ATTRVALUE
+_ATTRVALUE.fields_by_name['type'].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE
+_ATTRVALUE.fields_by_name['shape'].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO
+_ATTRVALUE.fields_by_name['tensor'].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO
+_ATTRVALUE.fields_by_name['list'].message_type = _ATTRVALUE_LISTVALUE
+_ATTRVALUE.fields_by_name['func'].message_type = _NAMEATTRLIST
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['s'])
+_ATTRVALUE.fields_by_name['s'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['i'])
+_ATTRVALUE.fields_by_name['i'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['f'])
+_ATTRVALUE.fields_by_name['f'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['b'])
+_ATTRVALUE.fields_by_name['b'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['type'])
+_ATTRVALUE.fields_by_name['type'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['shape'])
+_ATTRVALUE.fields_by_name['shape'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['tensor'])
+_ATTRVALUE.fields_by_name['tensor'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['list'])
+_ATTRVALUE.fields_by_name['list'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['func'])
+_ATTRVALUE.fields_by_name['func'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_ATTRVALUE.oneofs_by_name['value'].fields.append(
+  _ATTRVALUE.fields_by_name['placeholder'])
+_ATTRVALUE.fields_by_name['placeholder'].containing_oneof = _ATTRVALUE.oneofs_by_name['value']
+_NAMEATTRLIST_ATTRENTRY.fields_by_name['value'].message_type = _ATTRVALUE
+_NAMEATTRLIST_ATTRENTRY.containing_type = _NAMEATTRLIST
+_NAMEATTRLIST.fields_by_name['attr'].message_type = _NAMEATTRLIST_ATTRENTRY
+DESCRIPTOR.message_types_by_name['AttrValue'] = _ATTRVALUE
+DESCRIPTOR.message_types_by_name['NameAttrList'] = _NAMEATTRLIST
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+AttrValue = _reflection.GeneratedProtocolMessageType('AttrValue', (_message.Message,), dict(
+
+  ListValue = _reflection.GeneratedProtocolMessageType('ListValue', (_message.Message,), dict(
+    DESCRIPTOR = _ATTRVALUE_LISTVALUE,
+    __module__ = 'tensorboardX.src.attr_value_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.AttrValue.ListValue)
+    ))
+  ,
+  DESCRIPTOR = _ATTRVALUE,
+  __module__ = 'tensorboardX.src.attr_value_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.AttrValue)
+  ))
+_sym_db.RegisterMessage(AttrValue)
+_sym_db.RegisterMessage(AttrValue.ListValue)
+
+NameAttrList = _reflection.GeneratedProtocolMessageType('NameAttrList', (_message.Message,), dict(
+
+  AttrEntry = _reflection.GeneratedProtocolMessageType('AttrEntry', (_message.Message,), dict(
+    DESCRIPTOR = _NAMEATTRLIST_ATTRENTRY,
+    __module__ = 'tensorboardX.src.attr_value_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.NameAttrList.AttrEntry)
+    ))
+  ,
+  DESCRIPTOR = _NAMEATTRLIST,
+  __module__ = 'tensorboardX.src.attr_value_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.NameAttrList)
+  ))
+_sym_db.RegisterMessage(NameAttrList)
+_sym_db.RegisterMessage(NameAttrList.AttrEntry)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\017AttrValueProtosP\001\370\001\001'))
+_ATTRVALUE_LISTVALUE.fields_by_name['i'].has_options = True
+_ATTRVALUE_LISTVALUE.fields_by_name['i']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_ATTRVALUE_LISTVALUE.fields_by_name['f'].has_options = True
+_ATTRVALUE_LISTVALUE.fields_by_name['f']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_ATTRVALUE_LISTVALUE.fields_by_name['b'].has_options = True
+_ATTRVALUE_LISTVALUE.fields_by_name['b']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_ATTRVALUE_LISTVALUE.fields_by_name['type'].has_options = True
+_ATTRVALUE_LISTVALUE.fields_by_name['type']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_NAMEATTRLIST_ATTRENTRY.has_options = True
+_NAMEATTRLIST_ATTRENTRY._options = _descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/event_pb2.py

@@ -0,0 +1,376 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/event.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorboardX.src import summary_pb2 as tensorboardX_dot_src_dot_summary__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/event.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1ctensorboardX/src/event.proto\x12\x0btensorboard\x1a\x1etensorboardX/src/summary.proto\"\xbf\x02\n\x05\x45vent\x12\x11\n\twall_time\x18\x01 \x01(\x01\x12\x0c\n\x04step\x18\x02 \x01(\x03\x12\x16\n\x0c\x66ile_version\x18\x03 \x01(\tH\x00\x12\x13\n\tgraph_def\x18\x04 \x01(\x0cH\x00\x12\'\n\x07summary\x18\x05 \x01(\x0b\x32\x14.tensorboard.SummaryH\x00\x12.\n\x0blog_message\x18\x06 \x01(\x0b\x32\x17.tensorboard.LogMessageH\x00\x12.\n\x0bsession_log\x18\x07 \x01(\x0b\x32\x17.tensorboard.SessionLogH\x00\x12=\n\x13tagged_run_metadata\x18\x08 \x01(\x0b\x32\x1e.tensorboard.TaggedRunMetadataH\x00\x12\x18\n\x0emeta_graph_def\x18\t \x01(\x0cH\x00\x42\x06\n\x04what\"\x96\x01\n\nLogMessage\x12,\n\x05level\x18\x01 \x01(\x0e\x32\x1d.tensorboard.LogMessage.Level\x12\x0f\n\x07message\x18\x02 \x01(\t\"I\n\x05Level\x12\x0b\n\x07UNKNOWN\x10\x00\x12\t\n\x05\x44\x45\x42UG\x10\n\x12\x08\n\x04INFO\x10\x14\x12\x08\n\x04WARN\x10\x1e\x12\t\n\x05\x45RROR\x10(\x12\t\n\x05\x46\x41TAL\x10\x32\"\xb7\x01\n\nSessionLog\x12\x35\n\x06status\x18\x01 \x01(\x0e\x32%.tensorboard.SessionLog.SessionStatus\x12\x17\n\x0f\x63heckpoint_path\x18\x02 \x01(\t\x12\x0b\n\x03msg\x18\x03 \x01(\t\"L\n\rSessionStatus\x12\x16\n\x12STATUS_UNSPECIFIED\x10\x00\x12\t\n\x05START\x10\x01\x12\x08\n\x04STOP\x10\x02\x12\x0e\n\nCHECKPOINT\x10\x03\"6\n\x11TaggedRunMetadata\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x14\n\x0crun_metadata\x18\x02 \x01(\x0c\x42\'\n\x13org.tensorflow.utilB\x0b\x45ventProtosP\x01\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorboardX_dot_src_dot_summary__pb2.DESCRIPTOR,])
+
+
+
+_LOGMESSAGE_LEVEL = _descriptor.EnumDescriptor(
+  name='Level',
+  full_name='tensorboard.LogMessage.Level',
+  filename=None,
+  file=DESCRIPTOR,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='UNKNOWN', index=0, number=0,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DEBUG', index=1, number=10,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='INFO', index=2, number=20,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='WARN', index=3, number=30,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='ERROR', index=4, number=40,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='FATAL', index=5, number=50,
+      options=None,
+      type=None),
+  ],
+  containing_type=None,
+  options=None,
+  serialized_start=477,
+  serialized_end=550,
+)
+_sym_db.RegisterEnumDescriptor(_LOGMESSAGE_LEVEL)
+
+_SESSIONLOG_SESSIONSTATUS = _descriptor.EnumDescriptor(
+  name='SessionStatus',
+  full_name='tensorboard.SessionLog.SessionStatus',
+  filename=None,
+  file=DESCRIPTOR,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='STATUS_UNSPECIFIED', index=0, number=0,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='START', index=1, number=1,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='STOP', index=2, number=2,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='CHECKPOINT', index=3, number=3,
+      options=None,
+      type=None),
+  ],
+  containing_type=None,
+  options=None,
+  serialized_start=660,
+  serialized_end=736,
+)
+_sym_db.RegisterEnumDescriptor(_SESSIONLOG_SESSIONSTATUS)
+
+
+_EVENT = _descriptor.Descriptor(
+  name='Event',
+  full_name='tensorboard.Event',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='wall_time', full_name='tensorboard.Event.wall_time', index=0,
+      number=1, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='step', full_name='tensorboard.Event.step', index=1,
+      number=2, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='file_version', full_name='tensorboard.Event.file_version', index=2,
+      number=3, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='graph_def', full_name='tensorboard.Event.graph_def', index=3,
+      number=4, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='summary', full_name='tensorboard.Event.summary', index=4,
+      number=5, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='log_message', full_name='tensorboard.Event.log_message', index=5,
+      number=6, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='session_log', full_name='tensorboard.Event.session_log', index=6,
+      number=7, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tagged_run_metadata', full_name='tensorboard.Event.tagged_run_metadata', index=7,
+      number=8, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='meta_graph_def', full_name='tensorboard.Event.meta_graph_def', index=8,
+      number=9, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='what', full_name='tensorboard.Event.what',
+      index=0, containing_type=None, fields=[]),
+  ],
+  serialized_start=78,
+  serialized_end=397,
+)
+
+
+_LOGMESSAGE = _descriptor.Descriptor(
+  name='LogMessage',
+  full_name='tensorboard.LogMessage',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='level', full_name='tensorboard.LogMessage.level', index=0,
+      number=1, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='message', full_name='tensorboard.LogMessage.message', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _LOGMESSAGE_LEVEL,
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=400,
+  serialized_end=550,
+)
+
+
+_SESSIONLOG = _descriptor.Descriptor(
+  name='SessionLog',
+  full_name='tensorboard.SessionLog',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='status', full_name='tensorboard.SessionLog.status', index=0,
+      number=1, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='checkpoint_path', full_name='tensorboard.SessionLog.checkpoint_path', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='msg', full_name='tensorboard.SessionLog.msg', index=2,
+      number=3, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _SESSIONLOG_SESSIONSTATUS,
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=553,
+  serialized_end=736,
+)
+
+
+_TAGGEDRUNMETADATA = _descriptor.Descriptor(
+  name='TaggedRunMetadata',
+  full_name='tensorboard.TaggedRunMetadata',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='tag', full_name='tensorboard.TaggedRunMetadata.tag', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='run_metadata', full_name='tensorboard.TaggedRunMetadata.run_metadata', index=1,
+      number=2, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=738,
+  serialized_end=792,
+)
+
+_EVENT.fields_by_name['summary'].message_type = tensorboardX_dot_src_dot_summary__pb2._SUMMARY
+_EVENT.fields_by_name['log_message'].message_type = _LOGMESSAGE
+_EVENT.fields_by_name['session_log'].message_type = _SESSIONLOG
+_EVENT.fields_by_name['tagged_run_metadata'].message_type = _TAGGEDRUNMETADATA
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['file_version'])
+_EVENT.fields_by_name['file_version'].containing_oneof = _EVENT.oneofs_by_name['what']
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['graph_def'])
+_EVENT.fields_by_name['graph_def'].containing_oneof = _EVENT.oneofs_by_name['what']
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['summary'])
+_EVENT.fields_by_name['summary'].containing_oneof = _EVENT.oneofs_by_name['what']
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['log_message'])
+_EVENT.fields_by_name['log_message'].containing_oneof = _EVENT.oneofs_by_name['what']
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['session_log'])
+_EVENT.fields_by_name['session_log'].containing_oneof = _EVENT.oneofs_by_name['what']
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['tagged_run_metadata'])
+_EVENT.fields_by_name['tagged_run_metadata'].containing_oneof = _EVENT.oneofs_by_name['what']
+_EVENT.oneofs_by_name['what'].fields.append(
+  _EVENT.fields_by_name['meta_graph_def'])
+_EVENT.fields_by_name['meta_graph_def'].containing_oneof = _EVENT.oneofs_by_name['what']
+_LOGMESSAGE.fields_by_name['level'].enum_type = _LOGMESSAGE_LEVEL
+_LOGMESSAGE_LEVEL.containing_type = _LOGMESSAGE
+_SESSIONLOG.fields_by_name['status'].enum_type = _SESSIONLOG_SESSIONSTATUS
+_SESSIONLOG_SESSIONSTATUS.containing_type = _SESSIONLOG
+DESCRIPTOR.message_types_by_name['Event'] = _EVENT
+DESCRIPTOR.message_types_by_name['LogMessage'] = _LOGMESSAGE
+DESCRIPTOR.message_types_by_name['SessionLog'] = _SESSIONLOG
+DESCRIPTOR.message_types_by_name['TaggedRunMetadata'] = _TAGGEDRUNMETADATA
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+Event = _reflection.GeneratedProtocolMessageType('Event', (_message.Message,), dict(
+  DESCRIPTOR = _EVENT,
+  __module__ = 'tensorboardX.src.event_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.Event)
+  ))
+_sym_db.RegisterMessage(Event)
+
+LogMessage = _reflection.GeneratedProtocolMessageType('LogMessage', (_message.Message,), dict(
+  DESCRIPTOR = _LOGMESSAGE,
+  __module__ = 'tensorboardX.src.event_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.LogMessage)
+  ))
+_sym_db.RegisterMessage(LogMessage)
+
+SessionLog = _reflection.GeneratedProtocolMessageType('SessionLog', (_message.Message,), dict(
+  DESCRIPTOR = _SESSIONLOG,
+  __module__ = 'tensorboardX.src.event_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.SessionLog)
+  ))
+_sym_db.RegisterMessage(SessionLog)
+
+TaggedRunMetadata = _reflection.GeneratedProtocolMessageType('TaggedRunMetadata', (_message.Message,), dict(
+  DESCRIPTOR = _TAGGEDRUNMETADATA,
+  __module__ = 'tensorboardX.src.event_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.TaggedRunMetadata)
+  ))
+_sym_db.RegisterMessage(TaggedRunMetadata)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\023org.tensorflow.utilB\013EventProtosP\001\370\001\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/graph_pb2.py

@@ -0,0 +1,92 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/graph.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorboardX.src import node_def_pb2 as tensorboardX_dot_src_dot_node__def__pb2
+from tensorboardX.src import versions_pb2 as tensorboardX_dot_src_dot_versions__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/graph.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1ctensorboardX/src/graph.proto\x12\x0btensorboard\x1a\x1ftensorboardX/src/node_def.proto\x1a\x1ftensorboardX/src/versions.proto\"n\n\x08GraphDef\x12\"\n\x04node\x18\x01 \x03(\x0b\x32\x14.tensorboard.NodeDef\x12)\n\x08versions\x18\x04 \x01(\x0b\x32\x17.tensorboard.VersionDef\x12\x13\n\x07version\x18\x03 \x01(\x05\x42\x02\x18\x01\x42,\n\x18org.tensorflow.frameworkB\x0bGraphProtosP\x01\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorboardX_dot_src_dot_node__def__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_versions__pb2.DESCRIPTOR,])
+
+
+
+
+_GRAPHDEF = _descriptor.Descriptor(
+  name='GraphDef',
+  full_name='tensorboard.GraphDef',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='node', full_name='tensorboard.GraphDef.node', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='versions', full_name='tensorboard.GraphDef.versions', index=1,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='version', full_name='tensorboard.GraphDef.version', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\030\001'))),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=111,
+  serialized_end=221,
+)
+
+_GRAPHDEF.fields_by_name['node'].message_type = tensorboardX_dot_src_dot_node__def__pb2._NODEDEF
+_GRAPHDEF.fields_by_name['versions'].message_type = tensorboardX_dot_src_dot_versions__pb2._VERSIONDEF
+DESCRIPTOR.message_types_by_name['GraphDef'] = _GRAPHDEF
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+GraphDef = _reflection.GeneratedProtocolMessageType('GraphDef', (_message.Message,), dict(
+  DESCRIPTOR = _GRAPHDEF,
+  __module__ = 'tensorboardX.src.graph_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.GraphDef)
+  ))
+_sym_db.RegisterMessage(GraphDef)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\013GraphProtosP\001\370\001\001'))
+_GRAPHDEF.fields_by_name['version'].has_options = True
+_GRAPHDEF.fields_by_name['version']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\030\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/node_def_pb2.py

@@ -0,0 +1,151 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/node_def.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorboardX.src import attr_value_pb2 as tensorboardX_dot_src_dot_attr__value__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/node_def.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1ftensorboardX/src/node_def.proto\x12\x0btensorboard\x1a!tensorboardX/src/attr_value.proto\"\xb5\x01\n\x07NodeDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\n\n\x02op\x18\x02 \x01(\t\x12\r\n\x05input\x18\x03 \x03(\t\x12\x0e\n\x06\x64\x65vice\x18\x04 \x01(\t\x12,\n\x04\x61ttr\x18\x05 \x03(\x0b\x32\x1e.tensorboard.NodeDef.AttrEntry\x1a\x43\n\tAttrEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.tensorboard.AttrValue:\x02\x38\x01\x42*\n\x18org.tensorflow.frameworkB\tNodeProtoP\x01\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorboardX_dot_src_dot_attr__value__pb2.DESCRIPTOR,])
+
+
+
+
+_NODEDEF_ATTRENTRY = _descriptor.Descriptor(
+  name='AttrEntry',
+  full_name='tensorboard.NodeDef.AttrEntry',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='key', full_name='tensorboard.NodeDef.AttrEntry.key', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='value', full_name='tensorboard.NodeDef.AttrEntry.value', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=_descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001')),
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=198,
+  serialized_end=265,
+)
+
+_NODEDEF = _descriptor.Descriptor(
+  name='NodeDef',
+  full_name='tensorboard.NodeDef',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='name', full_name='tensorboard.NodeDef.name', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='op', full_name='tensorboard.NodeDef.op', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='input', full_name='tensorboard.NodeDef.input', index=2,
+      number=3, type=9, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='device', full_name='tensorboard.NodeDef.device', index=3,
+      number=4, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='attr', full_name='tensorboard.NodeDef.attr', index=4,
+      number=5, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[_NODEDEF_ATTRENTRY, ],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=84,
+  serialized_end=265,
+)
+
+_NODEDEF_ATTRENTRY.fields_by_name['value'].message_type = tensorboardX_dot_src_dot_attr__value__pb2._ATTRVALUE
+_NODEDEF_ATTRENTRY.containing_type = _NODEDEF
+_NODEDEF.fields_by_name['attr'].message_type = _NODEDEF_ATTRENTRY
+DESCRIPTOR.message_types_by_name['NodeDef'] = _NODEDEF
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+NodeDef = _reflection.GeneratedProtocolMessageType('NodeDef', (_message.Message,), dict(
+
+  AttrEntry = _reflection.GeneratedProtocolMessageType('AttrEntry', (_message.Message,), dict(
+    DESCRIPTOR = _NODEDEF_ATTRENTRY,
+    __module__ = 'tensorboardX.src.node_def_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.NodeDef.AttrEntry)
+    ))
+  ,
+  DESCRIPTOR = _NODEDEF,
+  __module__ = 'tensorboardX.src.node_def_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.NodeDef)
+  ))
+_sym_db.RegisterMessage(NodeDef)
+_sym_db.RegisterMessage(NodeDef.AttrEntry)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\tNodeProtoP\001\370\001\001'))
+_NODEDEF_ATTRENTRY.has_options = True
+_NODEDEF_ATTRENTRY._options = _descriptor._ParseOptions(descriptor_pb2.MessageOptions(), _b('8\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/plugin_pr_curve_pb2.py

@@ -0,0 +1,76 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/plugin_pr_curve.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/plugin_pr_curve.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n&tensorboardX/src/plugin_pr_curve.proto\x12\x0btensorboard\"<\n\x11PrCurvePluginData\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0enum_thresholds\x18\x02 \x01(\rb\x06proto3')
+)
+
+
+
+
+_PRCURVEPLUGINDATA = _descriptor.Descriptor(
+  name='PrCurvePluginData',
+  full_name='tensorboard.PrCurvePluginData',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='version', full_name='tensorboard.PrCurvePluginData.version', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='num_thresholds', full_name='tensorboard.PrCurvePluginData.num_thresholds', index=1,
+      number=2, type=13, cpp_type=3, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=55,
+  serialized_end=115,
+)
+
+DESCRIPTOR.message_types_by_name['PrCurvePluginData'] = _PRCURVEPLUGINDATA
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+PrCurvePluginData = _reflection.GeneratedProtocolMessageType('PrCurvePluginData', (_message.Message,), dict(
+  DESCRIPTOR = _PRCURVEPLUGINDATA,
+  __module__ = 'tensorboardX.src.plugin_pr_curve_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.PrCurvePluginData)
+  ))
+_sym_db.RegisterMessage(PrCurvePluginData)
+
+
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/resource_handle_pb2.py

@@ -0,0 +1,99 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/resource_handle.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/resource_handle.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n&tensorboardX/src/resource_handle.proto\x12\x0btensorboard\"r\n\x13ResourceHandleProto\x12\x0e\n\x06\x64\x65vice\x18\x01 \x01(\t\x12\x11\n\tcontainer\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x11\n\thash_code\x18\x04 \x01(\x04\x12\x17\n\x0fmaybe_type_name\x18\x05 \x01(\tB/\n\x18org.tensorflow.frameworkB\x0eResourceHandleP\x01\xf8\x01\x01\x62\x06proto3')
+)
+
+
+
+
+_RESOURCEHANDLEPROTO = _descriptor.Descriptor(
+  name='ResourceHandleProto',
+  full_name='tensorboard.ResourceHandleProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='device', full_name='tensorboard.ResourceHandleProto.device', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='container', full_name='tensorboard.ResourceHandleProto.container', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='name', full_name='tensorboard.ResourceHandleProto.name', index=2,
+      number=3, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='hash_code', full_name='tensorboard.ResourceHandleProto.hash_code', index=3,
+      number=4, type=4, cpp_type=4, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='maybe_type_name', full_name='tensorboard.ResourceHandleProto.maybe_type_name', index=4,
+      number=5, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=55,
+  serialized_end=169,
+)
+
+DESCRIPTOR.message_types_by_name['ResourceHandleProto'] = _RESOURCEHANDLEPROTO
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+ResourceHandleProto = _reflection.GeneratedProtocolMessageType('ResourceHandleProto', (_message.Message,), dict(
+  DESCRIPTOR = _RESOURCEHANDLEPROTO,
+  __module__ = 'tensorboardX.src.resource_handle_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.ResourceHandleProto)
+  ))
+_sym_db.RegisterMessage(ResourceHandleProto)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\016ResourceHandleP\001\370\001\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/summary_pb2.py

@@ -0,0 +1,532 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/summary.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorboardX.src import tensor_pb2 as tensorboardX_dot_src_dot_tensor__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/summary.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1etensorboardX/src/summary.proto\x12\x0btensorboard\x1a\x1dtensorboardX/src/tensor.proto\"\'\n\x12SummaryDescription\x12\x11\n\ttype_hint\x18\x01 \x01(\t\"\x87\x01\n\x0eHistogramProto\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03num\x18\x03 \x01(\x01\x12\x0b\n\x03sum\x18\x04 \x01(\x01\x12\x13\n\x0bsum_squares\x18\x05 \x01(\x01\x12\x18\n\x0c\x62ucket_limit\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x12\n\x06\x62ucket\x18\x07 \x03(\x01\x42\x02\x10\x01\"\x83\x01\n\x0fSummaryMetadata\x12<\n\x0bplugin_data\x18\x01 \x03(\x0b\x32\'.tensorboard.SummaryMetadata.PluginData\x1a\x32\n\nPluginData\x12\x13\n\x0bplugin_name\x18\x01 \x01(\t\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\"\xe4\x04\n\x07Summary\x12)\n\x05value\x18\x01 \x03(\x0b\x32\x1a.tensorboard.Summary.Value\x1aX\n\x05Image\x12\x0e\n\x06height\x18\x01 \x01(\x05\x12\r\n\x05width\x18\x02 \x01(\x05\x12\x12\n\ncolorspace\x18\x03 \x01(\x05\x12\x1c\n\x14\x65ncoded_image_string\x18\x04 \x01(\x0c\x1a}\n\x05\x41udio\x12\x13\n\x0bsample_rate\x18\x01 \x01(\x02\x12\x14\n\x0cnum_channels\x18\x02 \x01(\x03\x12\x15\n\rlength_frames\x18\x03 \x01(\x03\x12\x1c\n\x14\x65ncoded_audio_string\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\x1a\xd4\x02\n\x05Value\x12\x11\n\tnode_name\x18\x07 \x01(\t\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12.\n\x08metadata\x18\t \x01(\x0b\x32\x1c.tensorboard.SummaryMetadata\x12\x16\n\x0csimple_value\x18\x02 \x01(\x02H\x00\x12&\n\x1cobsolete_old_style_histogram\x18\x03 \x01(\x0cH\x00\x12+\n\x05image\x18\x04 \x01(\x0b\x32\x1a.tensorboard.Summary.ImageH\x00\x12,\n\x05histo\x18\x05 \x01(\x0b\x32\x1b.tensorboard.HistogramProtoH\x00\x12+\n\x05\x61udio\x18\x06 \x01(\x0b\x32\x1a.tensorboard.Summary.AudioH\x00\x12*\n\x06tensor\x18\x08 \x01(\x0b\x32\x18.tensorboard.TensorProtoH\x00\x42\x07\n\x05valueB.\n\x18org.tensorflow.frameworkB\rSummaryProtosP\x01\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorboardX_dot_src_dot_tensor__pb2.DESCRIPTOR,])
+
+
+
+
+_SUMMARYDESCRIPTION = _descriptor.Descriptor(
+  name='SummaryDescription',
+  full_name='tensorboard.SummaryDescription',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='type_hint', full_name='tensorboard.SummaryDescription.type_hint', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=78,
+  serialized_end=117,
+)
+
+
+_HISTOGRAMPROTO = _descriptor.Descriptor(
+  name='HistogramProto',
+  full_name='tensorboard.HistogramProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='min', full_name='tensorboard.HistogramProto.min', index=0,
+      number=1, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='max', full_name='tensorboard.HistogramProto.max', index=1,
+      number=2, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='num', full_name='tensorboard.HistogramProto.num', index=2,
+      number=3, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='sum', full_name='tensorboard.HistogramProto.sum', index=3,
+      number=4, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='sum_squares', full_name='tensorboard.HistogramProto.sum_squares', index=4,
+      number=5, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='bucket_limit', full_name='tensorboard.HistogramProto.bucket_limit', index=5,
+      number=6, type=1, cpp_type=5, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='bucket', full_name='tensorboard.HistogramProto.bucket', index=6,
+      number=7, type=1, cpp_type=5, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=120,
+  serialized_end=255,
+)
+
+
+_SUMMARYMETADATA_PLUGINDATA = _descriptor.Descriptor(
+  name='PluginData',
+  full_name='tensorboard.SummaryMetadata.PluginData',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='plugin_name', full_name='tensorboard.SummaryMetadata.PluginData.plugin_name', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='content', full_name='tensorboard.SummaryMetadata.PluginData.content', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=339,
+  serialized_end=389,
+)
+
+_SUMMARYMETADATA = _descriptor.Descriptor(
+  name='SummaryMetadata',
+  full_name='tensorboard.SummaryMetadata',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='plugin_data', full_name='tensorboard.SummaryMetadata.plugin_data', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[_SUMMARYMETADATA_PLUGINDATA, ],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=258,
+  serialized_end=389,
+)
+
+
+_SUMMARY_IMAGE = _descriptor.Descriptor(
+  name='Image',
+  full_name='tensorboard.Summary.Image',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='height', full_name='tensorboard.Summary.Image.height', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='width', full_name='tensorboard.Summary.Image.width', index=1,
+      number=2, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='colorspace', full_name='tensorboard.Summary.Image.colorspace', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='encoded_image_string', full_name='tensorboard.Summary.Image.encoded_image_string', index=3,
+      number=4, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=446,
+  serialized_end=534,
+)
+
+_SUMMARY_AUDIO = _descriptor.Descriptor(
+  name='Audio',
+  full_name='tensorboard.Summary.Audio',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='sample_rate', full_name='tensorboard.Summary.Audio.sample_rate', index=0,
+      number=1, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='num_channels', full_name='tensorboard.Summary.Audio.num_channels', index=1,
+      number=2, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='length_frames', full_name='tensorboard.Summary.Audio.length_frames', index=2,
+      number=3, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='encoded_audio_string', full_name='tensorboard.Summary.Audio.encoded_audio_string', index=3,
+      number=4, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='content_type', full_name='tensorboard.Summary.Audio.content_type', index=4,
+      number=5, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=536,
+  serialized_end=661,
+)
+
+_SUMMARY_VALUE = _descriptor.Descriptor(
+  name='Value',
+  full_name='tensorboard.Summary.Value',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='node_name', full_name='tensorboard.Summary.Value.node_name', index=0,
+      number=7, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tag', full_name='tensorboard.Summary.Value.tag', index=1,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='metadata', full_name='tensorboard.Summary.Value.metadata', index=2,
+      number=9, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='simple_value', full_name='tensorboard.Summary.Value.simple_value', index=3,
+      number=2, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='obsolete_old_style_histogram', full_name='tensorboard.Summary.Value.obsolete_old_style_histogram', index=4,
+      number=3, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='image', full_name='tensorboard.Summary.Value.image', index=5,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='histo', full_name='tensorboard.Summary.Value.histo', index=6,
+      number=5, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='audio', full_name='tensorboard.Summary.Value.audio', index=7,
+      number=6, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tensor', full_name='tensorboard.Summary.Value.tensor', index=8,
+      number=8, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='value', full_name='tensorboard.Summary.Value.value',
+      index=0, containing_type=None, fields=[]),
+  ],
+  serialized_start=664,
+  serialized_end=1004,
+)
+
+_SUMMARY = _descriptor.Descriptor(
+  name='Summary',
+  full_name='tensorboard.Summary',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='value', full_name='tensorboard.Summary.value', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[_SUMMARY_IMAGE, _SUMMARY_AUDIO, _SUMMARY_VALUE, ],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=392,
+  serialized_end=1004,
+)
+
+_SUMMARYMETADATA_PLUGINDATA.containing_type = _SUMMARYMETADATA
+_SUMMARYMETADATA.fields_by_name['plugin_data'].message_type = _SUMMARYMETADATA_PLUGINDATA
+_SUMMARY_IMAGE.containing_type = _SUMMARY
+_SUMMARY_AUDIO.containing_type = _SUMMARY
+_SUMMARY_VALUE.fields_by_name['metadata'].message_type = _SUMMARYMETADATA
+_SUMMARY_VALUE.fields_by_name['image'].message_type = _SUMMARY_IMAGE
+_SUMMARY_VALUE.fields_by_name['histo'].message_type = _HISTOGRAMPROTO
+_SUMMARY_VALUE.fields_by_name['audio'].message_type = _SUMMARY_AUDIO
+_SUMMARY_VALUE.fields_by_name['tensor'].message_type = tensorboardX_dot_src_dot_tensor__pb2._TENSORPROTO
+_SUMMARY_VALUE.containing_type = _SUMMARY
+_SUMMARY_VALUE.oneofs_by_name['value'].fields.append(
+  _SUMMARY_VALUE.fields_by_name['simple_value'])
+_SUMMARY_VALUE.fields_by_name['simple_value'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value']
+_SUMMARY_VALUE.oneofs_by_name['value'].fields.append(
+  _SUMMARY_VALUE.fields_by_name['obsolete_old_style_histogram'])
+_SUMMARY_VALUE.fields_by_name['obsolete_old_style_histogram'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value']
+_SUMMARY_VALUE.oneofs_by_name['value'].fields.append(
+  _SUMMARY_VALUE.fields_by_name['image'])
+_SUMMARY_VALUE.fields_by_name['image'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value']
+_SUMMARY_VALUE.oneofs_by_name['value'].fields.append(
+  _SUMMARY_VALUE.fields_by_name['histo'])
+_SUMMARY_VALUE.fields_by_name['histo'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value']
+_SUMMARY_VALUE.oneofs_by_name['value'].fields.append(
+  _SUMMARY_VALUE.fields_by_name['audio'])
+_SUMMARY_VALUE.fields_by_name['audio'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value']
+_SUMMARY_VALUE.oneofs_by_name['value'].fields.append(
+  _SUMMARY_VALUE.fields_by_name['tensor'])
+_SUMMARY_VALUE.fields_by_name['tensor'].containing_oneof = _SUMMARY_VALUE.oneofs_by_name['value']
+_SUMMARY.fields_by_name['value'].message_type = _SUMMARY_VALUE
+DESCRIPTOR.message_types_by_name['SummaryDescription'] = _SUMMARYDESCRIPTION
+DESCRIPTOR.message_types_by_name['HistogramProto'] = _HISTOGRAMPROTO
+DESCRIPTOR.message_types_by_name['SummaryMetadata'] = _SUMMARYMETADATA
+DESCRIPTOR.message_types_by_name['Summary'] = _SUMMARY
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+SummaryDescription = _reflection.GeneratedProtocolMessageType('SummaryDescription', (_message.Message,), dict(
+  DESCRIPTOR = _SUMMARYDESCRIPTION,
+  __module__ = 'tensorboardX.src.summary_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.SummaryDescription)
+  ))
+_sym_db.RegisterMessage(SummaryDescription)
+
+HistogramProto = _reflection.GeneratedProtocolMessageType('HistogramProto', (_message.Message,), dict(
+  DESCRIPTOR = _HISTOGRAMPROTO,
+  __module__ = 'tensorboardX.src.summary_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.HistogramProto)
+  ))
+_sym_db.RegisterMessage(HistogramProto)
+
+SummaryMetadata = _reflection.GeneratedProtocolMessageType('SummaryMetadata', (_message.Message,), dict(
+
+  PluginData = _reflection.GeneratedProtocolMessageType('PluginData', (_message.Message,), dict(
+    DESCRIPTOR = _SUMMARYMETADATA_PLUGINDATA,
+    __module__ = 'tensorboardX.src.summary_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.SummaryMetadata.PluginData)
+    ))
+  ,
+  DESCRIPTOR = _SUMMARYMETADATA,
+  __module__ = 'tensorboardX.src.summary_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.SummaryMetadata)
+  ))
+_sym_db.RegisterMessage(SummaryMetadata)
+_sym_db.RegisterMessage(SummaryMetadata.PluginData)
+
+Summary = _reflection.GeneratedProtocolMessageType('Summary', (_message.Message,), dict(
+
+  Image = _reflection.GeneratedProtocolMessageType('Image', (_message.Message,), dict(
+    DESCRIPTOR = _SUMMARY_IMAGE,
+    __module__ = 'tensorboardX.src.summary_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.Summary.Image)
+    ))
+  ,
+
+  Audio = _reflection.GeneratedProtocolMessageType('Audio', (_message.Message,), dict(
+    DESCRIPTOR = _SUMMARY_AUDIO,
+    __module__ = 'tensorboardX.src.summary_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.Summary.Audio)
+    ))
+  ,
+
+  Value = _reflection.GeneratedProtocolMessageType('Value', (_message.Message,), dict(
+    DESCRIPTOR = _SUMMARY_VALUE,
+    __module__ = 'tensorboardX.src.summary_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.Summary.Value)
+    ))
+  ,
+  DESCRIPTOR = _SUMMARY,
+  __module__ = 'tensorboardX.src.summary_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.Summary)
+  ))
+_sym_db.RegisterMessage(Summary)
+_sym_db.RegisterMessage(Summary.Image)
+_sym_db.RegisterMessage(Summary.Audio)
+_sym_db.RegisterMessage(Summary.Value)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\rSummaryProtosP\001\370\001\001'))
+_HISTOGRAMPROTO.fields_by_name['bucket_limit'].has_options = True
+_HISTOGRAMPROTO.fields_by_name['bucket_limit']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_HISTOGRAMPROTO.fields_by_name['bucket'].has_options = True
+_HISTOGRAMPROTO.fields_by_name['bucket']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/tensor_pb2.py

@@ -0,0 +1,185 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/tensor.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from tensorboardX.src import resource_handle_pb2 as tensorboardX_dot_src_dot_resource__handle__pb2
+from tensorboardX.src import tensor_shape_pb2 as tensorboardX_dot_src_dot_tensor__shape__pb2
+from tensorboardX.src import types_pb2 as tensorboardX_dot_src_dot_types__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/tensor.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1dtensorboardX/src/tensor.proto\x12\x0btensorboard\x1a&tensorboardX/src/resource_handle.proto\x1a#tensorboardX/src/tensor_shape.proto\x1a\x1ctensorboardX/src/types.proto\"\xa6\x03\n\x0bTensorProto\x12$\n\x05\x64type\x18\x01 \x01(\x0e\x32\x15.tensorboard.DataType\x12\x33\n\x0ctensor_shape\x18\x02 \x01(\x0b\x32\x1d.tensorboard.TensorShapeProto\x12\x16\n\x0eversion_number\x18\x03 \x01(\x05\x12\x16\n\x0etensor_content\x18\x04 \x01(\x0c\x12\x14\n\x08half_val\x18\r \x03(\x05\x42\x02\x10\x01\x12\x15\n\tfloat_val\x18\x05 \x03(\x02\x42\x02\x10\x01\x12\x16\n\ndouble_val\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x13\n\x07int_val\x18\x07 \x03(\x05\x42\x02\x10\x01\x12\x12\n\nstring_val\x18\x08 \x03(\x0c\x12\x18\n\x0cscomplex_val\x18\t \x03(\x02\x42\x02\x10\x01\x12\x15\n\tint64_val\x18\n \x03(\x03\x42\x02\x10\x01\x12\x14\n\x08\x62ool_val\x18\x0b \x03(\x08\x42\x02\x10\x01\x12\x18\n\x0c\x64\x63omplex_val\x18\x0c \x03(\x01\x42\x02\x10\x01\x12=\n\x13resource_handle_val\x18\x0e \x03(\x0b\x32 .tensorboard.ResourceHandleProtoB-\n\x18org.tensorflow.frameworkB\x0cTensorProtosP\x01\xf8\x01\x01\x62\x06proto3')
+  ,
+  dependencies=[tensorboardX_dot_src_dot_resource__handle__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_tensor__shape__pb2.DESCRIPTOR,tensorboardX_dot_src_dot_types__pb2.DESCRIPTOR,])
+
+
+
+
+_TENSORPROTO = _descriptor.Descriptor(
+  name='TensorProto',
+  full_name='tensorboard.TensorProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='dtype', full_name='tensorboard.TensorProto.dtype', index=0,
+      number=1, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tensor_shape', full_name='tensorboard.TensorProto.tensor_shape', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='version_number', full_name='tensorboard.TensorProto.version_number', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='tensor_content', full_name='tensorboard.TensorProto.tensor_content', index=3,
+      number=4, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b(""),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='half_val', full_name='tensorboard.TensorProto.half_val', index=4,
+      number=13, type=5, cpp_type=1, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='float_val', full_name='tensorboard.TensorProto.float_val', index=5,
+      number=5, type=2, cpp_type=6, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='double_val', full_name='tensorboard.TensorProto.double_val', index=6,
+      number=6, type=1, cpp_type=5, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='int_val', full_name='tensorboard.TensorProto.int_val', index=7,
+      number=7, type=5, cpp_type=1, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='string_val', full_name='tensorboard.TensorProto.string_val', index=8,
+      number=8, type=12, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='scomplex_val', full_name='tensorboard.TensorProto.scomplex_val', index=9,
+      number=9, type=2, cpp_type=6, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='int64_val', full_name='tensorboard.TensorProto.int64_val', index=10,
+      number=10, type=3, cpp_type=2, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='bool_val', full_name='tensorboard.TensorProto.bool_val', index=11,
+      number=11, type=8, cpp_type=7, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='dcomplex_val', full_name='tensorboard.TensorProto.dcomplex_val', index=12,
+      number=12, type=1, cpp_type=5, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=_descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))),
+    _descriptor.FieldDescriptor(
+      name='resource_handle_val', full_name='tensorboard.TensorProto.resource_handle_val', index=13,
+      number=14, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=154,
+  serialized_end=576,
+)
+
+_TENSORPROTO.fields_by_name['dtype'].enum_type = tensorboardX_dot_src_dot_types__pb2._DATATYPE
+_TENSORPROTO.fields_by_name['tensor_shape'].message_type = tensorboardX_dot_src_dot_tensor__shape__pb2._TENSORSHAPEPROTO
+_TENSORPROTO.fields_by_name['resource_handle_val'].message_type = tensorboardX_dot_src_dot_resource__handle__pb2._RESOURCEHANDLEPROTO
+DESCRIPTOR.message_types_by_name['TensorProto'] = _TENSORPROTO
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+TensorProto = _reflection.GeneratedProtocolMessageType('TensorProto', (_message.Message,), dict(
+  DESCRIPTOR = _TENSORPROTO,
+  __module__ = 'tensorboardX.src.tensor_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.TensorProto)
+  ))
+_sym_db.RegisterMessage(TensorProto)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\014TensorProtosP\001\370\001\001'))
+_TENSORPROTO.fields_by_name['half_val'].has_options = True
+_TENSORPROTO.fields_by_name['half_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['float_val'].has_options = True
+_TENSORPROTO.fields_by_name['float_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['double_val'].has_options = True
+_TENSORPROTO.fields_by_name['double_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['int_val'].has_options = True
+_TENSORPROTO.fields_by_name['int_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['scomplex_val'].has_options = True
+_TENSORPROTO.fields_by_name['scomplex_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['int64_val'].has_options = True
+_TENSORPROTO.fields_by_name['int64_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['bool_val'].has_options = True
+_TENSORPROTO.fields_by_name['bool_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+_TENSORPROTO.fields_by_name['dcomplex_val'].has_options = True
+_TENSORPROTO.fields_by_name['dcomplex_val']._options = _descriptor._ParseOptions(descriptor_pb2.FieldOptions(), _b('\020\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/tensor_shape_pb2.py

@@ -0,0 +1,125 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/tensor_shape.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/tensor_shape.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n#tensorboardX/src/tensor_shape.proto\x12\x0btensorboard\"{\n\x10TensorShapeProto\x12.\n\x03\x64im\x18\x02 \x03(\x0b\x32!.tensorboard.TensorShapeProto.Dim\x12\x14\n\x0cunknown_rank\x18\x03 \x01(\x08\x1a!\n\x03\x44im\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\tB2\n\x18org.tensorflow.frameworkB\x11TensorShapeProtosP\x01\xf8\x01\x01\x62\x06proto3')
+)
+
+
+
+
+_TENSORSHAPEPROTO_DIM = _descriptor.Descriptor(
+  name='Dim',
+  full_name='tensorboard.TensorShapeProto.Dim',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='size', full_name='tensorboard.TensorShapeProto.Dim.size', index=0,
+      number=1, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='name', full_name='tensorboard.TensorShapeProto.Dim.name', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=_b("").decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=142,
+  serialized_end=175,
+)
+
+_TENSORSHAPEPROTO = _descriptor.Descriptor(
+  name='TensorShapeProto',
+  full_name='tensorboard.TensorShapeProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='dim', full_name='tensorboard.TensorShapeProto.dim', index=0,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='unknown_rank', full_name='tensorboard.TensorShapeProto.unknown_rank', index=1,
+      number=3, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[_TENSORSHAPEPROTO_DIM, ],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=52,
+  serialized_end=175,
+)
+
+_TENSORSHAPEPROTO_DIM.containing_type = _TENSORSHAPEPROTO
+_TENSORSHAPEPROTO.fields_by_name['dim'].message_type = _TENSORSHAPEPROTO_DIM
+DESCRIPTOR.message_types_by_name['TensorShapeProto'] = _TENSORSHAPEPROTO
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+TensorShapeProto = _reflection.GeneratedProtocolMessageType('TensorShapeProto', (_message.Message,), dict(
+
+  Dim = _reflection.GeneratedProtocolMessageType('Dim', (_message.Message,), dict(
+    DESCRIPTOR = _TENSORSHAPEPROTO_DIM,
+    __module__ = 'tensorboardX.src.tensor_shape_pb2'
+    # @@protoc_insertion_point(class_scope:tensorboard.TensorShapeProto.Dim)
+    ))
+  ,
+  DESCRIPTOR = _TENSORSHAPEPROTO,
+  __module__ = 'tensorboardX.src.tensor_shape_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.TensorShapeProto)
+  ))
+_sym_db.RegisterMessage(TensorShapeProto)
+_sym_db.RegisterMessage(TensorShapeProto.Dim)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\021TensorShapeProtosP\001\370\001\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/types_pb2.py

@@ -0,0 +1,254 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/types.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf.internal import enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/types.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1ctensorboardX/src/types.proto\x12\x0btensorboard*\xc2\x05\n\x08\x44\x61taType\x12\x0e\n\nDT_INVALID\x10\x00\x12\x0c\n\x08\x44T_FLOAT\x10\x01\x12\r\n\tDT_DOUBLE\x10\x02\x12\x0c\n\x08\x44T_INT32\x10\x03\x12\x0c\n\x08\x44T_UINT8\x10\x04\x12\x0c\n\x08\x44T_INT16\x10\x05\x12\x0b\n\x07\x44T_INT8\x10\x06\x12\r\n\tDT_STRING\x10\x07\x12\x10\n\x0c\x44T_COMPLEX64\x10\x08\x12\x0c\n\x08\x44T_INT64\x10\t\x12\x0b\n\x07\x44T_BOOL\x10\n\x12\x0c\n\x08\x44T_QINT8\x10\x0b\x12\r\n\tDT_QUINT8\x10\x0c\x12\r\n\tDT_QINT32\x10\r\x12\x0f\n\x0b\x44T_BFLOAT16\x10\x0e\x12\r\n\tDT_QINT16\x10\x0f\x12\x0e\n\nDT_QUINT16\x10\x10\x12\r\n\tDT_UINT16\x10\x11\x12\x11\n\rDT_COMPLEX128\x10\x12\x12\x0b\n\x07\x44T_HALF\x10\x13\x12\x0f\n\x0b\x44T_RESOURCE\x10\x14\x12\x10\n\x0c\x44T_FLOAT_REF\x10\x65\x12\x11\n\rDT_DOUBLE_REF\x10\x66\x12\x10\n\x0c\x44T_INT32_REF\x10g\x12\x10\n\x0c\x44T_UINT8_REF\x10h\x12\x10\n\x0c\x44T_INT16_REF\x10i\x12\x0f\n\x0b\x44T_INT8_REF\x10j\x12\x11\n\rDT_STRING_REF\x10k\x12\x14\n\x10\x44T_COMPLEX64_REF\x10l\x12\x10\n\x0c\x44T_INT64_REF\x10m\x12\x0f\n\x0b\x44T_BOOL_REF\x10n\x12\x10\n\x0c\x44T_QINT8_REF\x10o\x12\x11\n\rDT_QUINT8_REF\x10p\x12\x11\n\rDT_QINT32_REF\x10q\x12\x13\n\x0f\x44T_BFLOAT16_REF\x10r\x12\x11\n\rDT_QINT16_REF\x10s\x12\x12\n\x0e\x44T_QUINT16_REF\x10t\x12\x11\n\rDT_UINT16_REF\x10u\x12\x15\n\x11\x44T_COMPLEX128_REF\x10v\x12\x0f\n\x0b\x44T_HALF_REF\x10w\x12\x13\n\x0f\x44T_RESOURCE_REF\x10xB,\n\x18org.tensorflow.frameworkB\x0bTypesProtosP\x01\xf8\x01\x01\x62\x06proto3')
+)
+
+_DATATYPE = _descriptor.EnumDescriptor(
+  name='DataType',
+  full_name='tensorboard.DataType',
+  filename=None,
+  file=DESCRIPTOR,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='DT_INVALID', index=0, number=0,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_FLOAT', index=1, number=1,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_DOUBLE', index=2, number=2,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT32', index=3, number=3,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_UINT8', index=4, number=4,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT16', index=5, number=5,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT8', index=6, number=6,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_STRING', index=7, number=7,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_COMPLEX64', index=8, number=8,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT64', index=9, number=9,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_BOOL', index=10, number=10,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QINT8', index=11, number=11,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QUINT8', index=12, number=12,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QINT32', index=13, number=13,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_BFLOAT16', index=14, number=14,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QINT16', index=15, number=15,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QUINT16', index=16, number=16,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_UINT16', index=17, number=17,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_COMPLEX128', index=18, number=18,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_HALF', index=19, number=19,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_RESOURCE', index=20, number=20,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_FLOAT_REF', index=21, number=101,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_DOUBLE_REF', index=22, number=102,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT32_REF', index=23, number=103,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_UINT8_REF', index=24, number=104,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT16_REF', index=25, number=105,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT8_REF', index=26, number=106,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_STRING_REF', index=27, number=107,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_COMPLEX64_REF', index=28, number=108,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT64_REF', index=29, number=109,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_BOOL_REF', index=30, number=110,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QINT8_REF', index=31, number=111,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QUINT8_REF', index=32, number=112,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QINT32_REF', index=33, number=113,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_BFLOAT16_REF', index=34, number=114,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QINT16_REF', index=35, number=115,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_QUINT16_REF', index=36, number=116,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_UINT16_REF', index=37, number=117,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_COMPLEX128_REF', index=38, number=118,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_HALF_REF', index=39, number=119,
+      options=None,
+      type=None),
+    _descriptor.EnumValueDescriptor(
+      name='DT_RESOURCE_REF', index=40, number=120,
+      options=None,
+      type=None),
+  ],
+  containing_type=None,
+  options=None,
+  serialized_start=46,
+  serialized_end=752,
+)
+_sym_db.RegisterEnumDescriptor(_DATATYPE)
+
+DataType = enum_type_wrapper.EnumTypeWrapper(_DATATYPE)
+DT_INVALID = 0
+DT_FLOAT = 1
+DT_DOUBLE = 2
+DT_INT32 = 3
+DT_UINT8 = 4
+DT_INT16 = 5
+DT_INT8 = 6
+DT_STRING = 7
+DT_COMPLEX64 = 8
+DT_INT64 = 9
+DT_BOOL = 10
+DT_QINT8 = 11
+DT_QUINT8 = 12
+DT_QINT32 = 13
+DT_BFLOAT16 = 14
+DT_QINT16 = 15
+DT_QUINT16 = 16
+DT_UINT16 = 17
+DT_COMPLEX128 = 18
+DT_HALF = 19
+DT_RESOURCE = 20
+DT_FLOAT_REF = 101
+DT_DOUBLE_REF = 102
+DT_INT32_REF = 103
+DT_UINT8_REF = 104
+DT_INT16_REF = 105
+DT_INT8_REF = 106
+DT_STRING_REF = 107
+DT_COMPLEX64_REF = 108
+DT_INT64_REF = 109
+DT_BOOL_REF = 110
+DT_QINT8_REF = 111
+DT_QUINT8_REF = 112
+DT_QINT32_REF = 113
+DT_BFLOAT16_REF = 114
+DT_QINT16_REF = 115
+DT_QUINT16_REF = 116
+DT_UINT16_REF = 117
+DT_COMPLEX128_REF = 118
+DT_HALF_REF = 119
+DT_RESOURCE_REF = 120
+
+
+DESCRIPTOR.enum_types_by_name['DataType'] = _DATATYPE
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\013TypesProtosP\001\370\001\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/src/versions_pb2.py

@@ -0,0 +1,85 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: tensorboardX/src/versions.proto
+
+import sys
+_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import descriptor_pb2
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='tensorboardX/src/versions.proto',
+  package='tensorboard',
+  syntax='proto3',
+  serialized_pb=_b('\n\x1ftensorboardX/src/versions.proto\x12\x0btensorboard\"K\n\nVersionDef\x12\x10\n\x08producer\x18\x01 \x01(\x05\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\x05\x12\x15\n\rbad_consumers\x18\x03 \x03(\x05\x42/\n\x18org.tensorflow.frameworkB\x0eVersionsProtosP\x01\xf8\x01\x01\x62\x06proto3')
+)
+
+
+
+
+_VERSIONDEF = _descriptor.Descriptor(
+  name='VersionDef',
+  full_name='tensorboard.VersionDef',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='producer', full_name='tensorboard.VersionDef.producer', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='min_consumer', full_name='tensorboard.VersionDef.min_consumer', index=1,
+      number=2, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+    _descriptor.FieldDescriptor(
+      name='bad_consumers', full_name='tensorboard.VersionDef.bad_consumers', index=2,
+      number=3, type=5, cpp_type=1, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      options=None),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  options=None,
+  is_extendable=False,
+  syntax='proto3',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=48,
+  serialized_end=123,
+)
+
+DESCRIPTOR.message_types_by_name['VersionDef'] = _VERSIONDEF
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+VersionDef = _reflection.GeneratedProtocolMessageType('VersionDef', (_message.Message,), dict(
+  DESCRIPTOR = _VERSIONDEF,
+  __module__ = 'tensorboardX.src.versions_pb2'
+  # @@protoc_insertion_point(class_scope:tensorboard.VersionDef)
+  ))
+_sym_db.RegisterMessage(VersionDef)
+
+
+DESCRIPTOR.has_options = True
+DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), _b('\n\030org.tensorflow.frameworkB\016VersionsProtosP\001\370\001\001'))
+# @@protoc_insertion_point(module_scope)

tensorboardX/summary.py

@@ -0,0 +1,264 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""## Generation of summaries.
+### Class for writing Summaries
+@@FileWriter
+@@FileWriterCache
+### Summary Ops
+@@tensor_summary
+@@scalar
+@@histogram
+@@audio
+@@image
+@@merge
+@@merge_all
+## Utilities
+@@get_summary_description
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import re as _re
+import bisect
+from six import StringIO
+from six.moves import range
+import numpy as np
+# pylint: disable=unused-import
+from .src.summary_pb2 import Summary
+from .src.summary_pb2 import HistogramProto
+from .src.summary_pb2 import SummaryMetadata
+from .src.tensor_pb2 import TensorProto
+from .src.tensor_shape_pb2 import TensorShapeProto
+from .src.plugin_pr_curve_pb2 import PrCurvePluginData
+from .x2num import makenp
+
+_INVALID_TAG_CHARACTERS = _re.compile(r'[^-/\w\.]')
+
+
+def _clean_tag(name):
+    # In the past, the first argument to summary ops was a tag, which allowed
+    # arbitrary characters. Now we are changing the first argument to be the node
+    # name. This has a number of advantages (users of summary ops now can
+    # take advantage of the tf name scope system) but risks breaking existing
+    # usage, because a much smaller set of characters are allowed in node names.
+    # This function replaces all illegal characters with _s, and logs a warning.
+    # It also strips leading slashes from the name.
+    if name is not None:
+        new_name = _INVALID_TAG_CHARACTERS.sub('_', name)
+        new_name = new_name.lstrip('/')  # Remove leading slashes
+        if new_name != name:
+            logging.info('Summary name %s is illegal; using %s instead.' % (name, new_name))
+            name = new_name
+    return name
+
+
+def scalar(name, scalar, collections=None):
+    """Outputs a `Summary` protocol buffer containing a single scalar value.
+    The generated Summary has a Tensor.proto containing the input Tensor.
+    Args:
+      name: A name for the generated node. Will also serve as the series name in
+        TensorBoard.
+      tensor: A real numeric Tensor containing a single value.
+      collections: Optional list of graph collections keys. The new summary op is
+        added to these collections. Defaults to `[GraphKeys.SUMMARIES]`.
+    Returns:
+      A scalar `Tensor` of type `string`. Which contains a `Summary` protobuf.
+    Raises:
+      ValueError: If tensor has the wrong shape or type.
+    """
+    name = _clean_tag(name)
+    scalar = makenp(scalar)
+    assert(scalar.squeeze().ndim == 0), 'scalar should be 0D'
+    scalar = float(scalar)
+    return Summary(value=[Summary.Value(tag=name, simple_value=scalar)])
+
+
+def histogram(name, values, bins, collections=None):
+    # pylint: disable=line-too-long
+    """Outputs a `Summary` protocol buffer with a histogram.
+    The generated
+    [`Summary`](https://www.tensorflow.org/code/tensorflow/core/framework/summary.proto)
+    has one summary value containing a histogram for `values`.
+    This op reports an `InvalidArgument` error if any value is not finite.
+    Args:
+      name: A name for the generated node. Will also serve as a series name in
+        TensorBoard.
+      values: A real numeric `Tensor`. Any shape. Values to use to
+        build the histogram.
+      collections: Optional list of graph collections keys. The new summary op is
+        added to these collections. Defaults to `[GraphKeys.SUMMARIES]`.
+    Returns:
+      A scalar `Tensor` of type `string`. The serialized `Summary` protocol
+      buffer.
+    """
+    name = _clean_tag(name)
+    values = makenp(values)
+    hist = make_histogram(values.astype(float), bins)
+    return Summary(value=[Summary.Value(tag=name, histo=hist)])
+
+
+def make_histogram(values, bins):
+    """Convert values into a histogram proto using logic from histogram.cc."""
+    values = values.reshape(-1)
+    counts, limits = np.histogram(values, bins=bins)
+    limits = limits[1:]
+
+    sum_sq = values.dot(values)
+    return HistogramProto(min=values.min(),
+                          max=values.max(),
+                          num=len(values),
+                          sum=values.sum(),
+                          sum_squares=sum_sq,
+                          bucket_limit=limits,
+                          bucket=counts)
+
+
+def image(tag, tensor):
+    """Outputs a `Summary` protocol buffer with images.
+    The summary has up to `max_images` summary values containing images. The
+    images are built from `tensor` which must be 3-D with shape `[height, width,
+    channels]` and where `channels` can be:
+    *  1: `tensor` is interpreted as Grayscale.
+    *  3: `tensor` is interpreted as RGB.
+    *  4: `tensor` is interpreted as RGBA.
+    The `name` in the outputted Summary.Value protobufs is generated based on the
+    name, with a suffix depending on the max_outputs setting:
+    *  If `max_outputs` is 1, the summary value tag is '*name*/image'.
+    *  If `max_outputs` is greater than 1, the summary value tags are
+       generated sequentially as '*name*/image/0', '*name*/image/1', etc.
+    Args:
+      tag: A name for the generated node. Will also serve as a series name in
+        TensorBoard.
+      tensor: A 3-D `uint8` or `float32` `Tensor` of shape `[height, width,
+        channels]` where `channels` is 1, 3, or 4.
+    Returns:
+      A scalar `Tensor` of type `string`. The serialized `Summary` protocol
+      buffer.
+    """
+    tag = _clean_tag(tag)
+    tensor = makenp(tensor, 'IMG')
+    tensor = tensor.astype(np.float32)
+    tensor = (tensor * 255).astype(np.uint8)
+    image = make_image(tensor)
+    return Summary(value=[Summary.Value(tag=tag, image=image)])
+
+
+def make_image(tensor):
+    """Convert an numpy representation image to Image protobuf"""
+    from PIL import Image
+    height, width, channel = tensor.shape
+    image = Image.fromarray(tensor)
+    import io
+    output = io.BytesIO()
+    image.save(output, format='PNG')
+    image_string = output.getvalue()
+    output.close()
+    return Summary.Image(height=height,
+                         width=width,
+                         colorspace=channel,
+                         encoded_image_string=image_string)
+
+
+def audio(tag, tensor, sample_rate=44100):
+    tensor = makenp(tensor)
+    tensor = tensor.squeeze()
+    if abs(tensor).max() > 1:
+        print('warning: audio amplitude out of range, auto clipped.')
+        tensor = tensor.clip(-1, 1)
+    assert(tensor.ndim == 1), 'input tensor should be 1 dimensional.'
+
+    tensor_list = [int(32767.0 * x) for x in tensor]
+    import io
+    import wave
+    import struct
+    fio = io.BytesIO()
+    Wave_write = wave.open(fio, 'wb')
+    Wave_write.setnchannels(1)
+    Wave_write.setsampwidth(2)
+    Wave_write.setframerate(sample_rate)
+    tensor_enc = b''
+    for v in tensor_list:
+        tensor_enc += struct.pack('<h', v)
+
+    Wave_write.writeframes(tensor_enc)
+    Wave_write.close()
+    audio_string = fio.getvalue()
+    fio.close()
+    audio = Summary.Audio(sample_rate=sample_rate,
+                          num_channels=1,
+                          length_frames=len(tensor_list),
+                          encoded_audio_string=audio_string,
+                          content_type='audio/wav')
+    return Summary(value=[Summary.Value(tag=tag, audio=audio)])
+
+
+def text(tag, text):
+    import json
+    PluginData = [SummaryMetadata.PluginData(plugin_name='text')]
+    smd = SummaryMetadata(plugin_data=PluginData)
+    tensor = TensorProto(dtype='DT_STRING',
+                         string_val=[text.encode(encoding='utf_8')],
+                         tensor_shape=TensorShapeProto(dim=[TensorShapeProto.Dim(size=1)]))
+    return Summary(value=[Summary.Value(node_name=tag, metadata=smd, tensor=tensor)])
+
+
+def pr_curve(tag, labels, predictions, num_thresholds=127, weights=None):
+    if num_thresholds > 127:  # wierd, value > 127 breaks protobuf
+        num_thresholds = 127
+    data = compute_curve(labels, predictions, num_thresholds=num_thresholds, weights=weights)
+    pr_curve_plugin_data = PrCurvePluginData(version=0, num_thresholds=num_thresholds).SerializeToString()
+    PluginData = [SummaryMetadata.PluginData(plugin_name='pr_curves', content=pr_curve_plugin_data)]
+    smd = SummaryMetadata(plugin_data=PluginData)
+    tensor = TensorProto(dtype='DT_FLOAT',
+                         float_val=data.reshape(-1).tolist(),
+                         tensor_shape=TensorShapeProto(
+                             dim=[TensorShapeProto.Dim(size=data.shape[0]), TensorShapeProto.Dim(size=data.shape[1])]))
+    return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)])
+
+
+# https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/pr_curve/summary.py
+def compute_curve(labels, predictions, num_thresholds=None, weights=None):
+    _MINIMUM_COUNT = 1e-7
+
+    if weights is None:
+        weights = 1.0
+
+    # Compute bins of true positives and false positives.
+    bucket_indices = np.int32(np.floor(predictions * (num_thresholds - 1)))
+    float_labels = labels.astype(np.float)
+    histogram_range = (0, num_thresholds - 1)
+    tp_buckets, _ = np.histogram(
+        bucket_indices,
+        bins=num_thresholds,
+        range=histogram_range,
+        weights=float_labels * weights)
+    fp_buckets, _ = np.histogram(
+        bucket_indices,
+        bins=num_thresholds,
+        range=histogram_range,
+        weights=(1.0 - float_labels) * weights)
+
+    # Obtain the reverse cumulative sum.
+    tp = np.cumsum(tp_buckets[::-1])[::-1]
+    fp = np.cumsum(fp_buckets[::-1])[::-1]
+    tn = fp[0] - fp
+    fn = tp[0] - tp
+    precision = tp / np.maximum(_MINIMUM_COUNT, tp + fp)
+    recall = tp / np.maximum(_MINIMUM_COUNT, tp + fn)
+    return np.stack((tp, fp, tn, fn, precision, recall))

tensorboardX/writer.py

@@ -0,0 +1,485 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Provides an API for generating Event protocol buffers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+import json
+import os
+from .src import event_pb2
+from .src import summary_pb2
+from .src import graph_pb2
+from .event_file_writer import EventFileWriter
+from .summary import scalar, histogram, image, audio, text, pr_curve
+from .graph import graph
+from .graph_onnx import gg
+from .embedding import make_mat, make_sprite, make_tsv, append_pbtxt
+
+
+class SummaryToEventTransformer(object):
+    """Abstractly implements the SummaryWriter API.
+    This API basically implements a number of endpoints (add_summary,
+    add_session_log, etc). The endpoints all generate an event protobuf, which is
+    passed to the contained event_writer.
+    @@__init__
+    @@add_summary
+    @@add_session_log
+    @@add_graph
+    @@add_meta_graph
+    @@add_run_metadata
+    """
+
+    def __init__(self, event_writer, graph=None, graph_def=None):
+        """Creates a `SummaryWriter` and an event file.
+        On construction the summary writer creates a new event file in `logdir`.
+        This event file will contain `Event` protocol buffers constructed when you
+        call one of the following functions: `add_summary()`, `add_session_log()`,
+        `add_event()`, or `add_graph()`.
+        If you pass a `Graph` to the constructor it is added to
+        the event file. (This is equivalent to calling `add_graph()` later).
+        TensorBoard will pick the graph from the file and display it graphically so
+        you can interactively explore the graph you built. You will usually pass
+        the graph from the session in which you launched it:
+        ```python
+        ...create a graph...
+        # Launch the graph in a session.
+        sess = tf.Session()
+        # Create a summary writer, add the 'graph' to the event file.
+        writer = tf.summary.FileWriter(<some-directory>, sess.graph)
+        ```
+        Args:
+          event_writer: An EventWriter. Implements add_event method.
+          graph: A `Graph` object, such as `sess.graph`.
+          graph_def: DEPRECATED: Use the `graph` argument instead.
+        """
+        self.event_writer = event_writer
+        # For storing used tags for session.run() outputs.
+        self._session_run_tags = {}
+        # TODO(zihaolucky). pass this an empty graph to check whether it's necessary.
+        # currently we don't support graph in MXNet using tensorboard.
+
+    def add_summary(self, summary, global_step=None):
+        """Adds a `Summary` protocol buffer to the event file.
+        This method wraps the provided summary in an `Event` protocol buffer
+        and adds it to the event file.
+        You can pass the result of evaluating any summary op, using
+        [`Session.run()`](client.md#Session.run) or
+        [`Tensor.eval()`](framework.md#Tensor.eval), to this
+        function. Alternatively, you can pass a `tf.Summary` protocol
+        buffer that you populate with your own data. The latter is
+        commonly done to report evaluation results in event files.
+        Args:
+          summary: A `Summary` protocol buffer, optionally serialized as a string.
+          global_step: Number. Optional global step value to record with the
+            summary.
+        """
+        if isinstance(summary, bytes):
+            summ = summary_pb2.Summary()
+            summ.ParseFromString(summary)
+            summary = summ
+        event = event_pb2.Event(summary=summary)
+        self._add_event(event, global_step)
+
+    def add_graph_onnx(self, graph):
+        """Adds a `Graph` protocol buffer to the event file.
+        """
+        event = event_pb2.Event(graph_def=graph.SerializeToString())
+        self._add_event(event, None)
+
+    def add_graph(self, graph):
+        """Adds a `Graph` protocol buffer to the event file.
+        """
+        event = event_pb2.Event(graph_def=graph.SerializeToString())
+        self._add_event(event, None)
+
+    def add_session_log(self, session_log, global_step=None):
+        """Adds a `SessionLog` protocol buffer to the event file.
+        This method wraps the provided session in an `Event` protocol buffer
+        and adds it to the event file.
+        Args:
+          session_log: A `SessionLog` protocol buffer.
+          global_step: Number. Optional global step value to record with the
+            summary.
+        """
+        event = event_pb2.Event(session_log=session_log)
+        self._add_event(event, global_step)
+
+    def _add_event(self, event, step):
+        event.wall_time = time.time()
+        if step is not None:
+            event.step = int(step)
+        self.event_writer.add_event(event)
+
+
+class FileWriter(SummaryToEventTransformer):
+    """Writes `Summary` protocol buffers to event files.
+    The `FileWriter` class provides a mechanism to create an event file in a
+    given directory and add summaries and events to it. The class updates the
+    file contents asynchronously. This allows a training program to call methods
+    to add data to the file directly from the training loop, without slowing down
+    training.
+    @@__init__
+    @@add_summary
+    @@add_session_log
+    @@add_event
+    @@add_graph
+    @@add_run_metadata
+    @@get_logdir
+    @@flush
+    @@close
+    """
+
+    def __init__(self,
+                 logdir,
+                 graph=None,
+                 max_queue=10,
+                 flush_secs=120,
+                 graph_def=None):
+        """Creates a `FileWriter` and an event file.
+        On construction the summary writer creates a new event file in `logdir`.
+        This event file will contain `Event` protocol buffers constructed when you
+        call one of the following functions: `add_summary()`, `add_session_log()`,
+        `add_event()`, or `add_graph()`.
+        If you pass a `Graph` to the constructor it is added to
+        the event file. (This is equivalent to calling `add_graph()` later).
+        TensorBoard will pick the graph from the file and display it graphically so
+        you can interactively explore the graph you built. You will usually pass
+        the graph from the session in which you launched it:
+        ```python
+        ...create a graph...
+        # Launch the graph in a session.
+        sess = tf.Session()
+        # Create a summary writer, add the 'graph' to the event file.
+        writer = tf.summary.FileWriter(<some-directory>, sess.graph)
+        ```
+        The other arguments to the constructor control the asynchronous writes to
+        the event file:
+        *  `flush_secs`: How often, in seconds, to flush the added summaries
+           and events to disk.
+        *  `max_queue`: Maximum number of summaries or events pending to be
+           written to disk before one of the 'add' calls block.
+        Args:
+          logdir: A string. Directory where event file will be written.
+          graph: A `Graph` object, such as `sess.graph`.
+          max_queue: Integer. Size of the queue for pending events and summaries.
+          flush_secs: Number. How often, in seconds, to flush the
+            pending events and summaries to disk.
+          graph_def: DEPRECATED: Use the `graph` argument instead.
+        """
+        event_writer = EventFileWriter(logdir, max_queue, flush_secs)
+        super(FileWriter, self).__init__(event_writer, graph, graph_def)
+
+    def get_logdir(self):
+        """Returns the directory where event file will be written."""
+        return self.event_writer.get_logdir()
+
+    def add_event(self, event):
+        """Adds an event to the event file.
+        Args:
+          event: An `Event` protocol buffer.
+        """
+        self.event_writer.add_event(event)
+
+    def flush(self):
+        """Flushes the event file to disk.
+        Call this method to make sure that all pending events have been written to
+        disk.
+        """
+        self.event_writer.flush()
+
+    def close(self):
+        """Flushes the event file to disk and close the file.
+        Call this method when you do not need the summary writer anymore.
+        """
+        self.event_writer.close()
+
+    def reopen(self):
+        """Reopens the EventFileWriter.
+        Can be called after `close()` to add more events in the same directory.
+        The events will go into a new events file.
+        Does nothing if the EventFileWriter was not closed.
+        """
+        self.event_writer.reopen()
+
+
+class SummaryWriter(object):
+    """Writes `Summary` directly to event files.
+    The `SummaryWriter` class provides a high-level api to create an event file in a
+    given directory and add summaries and events to it. The class updates the
+    file contents asynchronously. This allows a training program to call methods
+    to add data to the file directly from the training loop, without slowing down
+    training.
+    """
+    def __init__(self, log_dir=None, comment=''):
+        """
+        Args:
+            log_dir (string): save location, default is: runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each
+              run. Use hierarchical folder structure to compare between runs easily. e.g. 'runs/exp1', 'runs/exp2'
+            comment (string): comment that appends to the default log_dir
+        """
+        if not log_dir:
+            import socket
+            from datetime import datetime
+            current_time = datetime.now().strftime('%b%d_%H-%M-%S')
+            log_dir = os.path.join('runs', current_time + '_' + socket.gethostname() + comment)
+        self.file_writer = FileWriter(logdir=log_dir)
+        v = 1E-12
+        buckets = []
+        neg_buckets = []
+        while v < 1E20:
+            buckets.append(v)
+            neg_buckets.append(-v)
+            v *= 1.1
+        self.default_bins = neg_buckets[::-1] + [0] + buckets
+        self.text_tags = []
+        #
+        self.all_writers = {self.file_writer.get_logdir(): self.file_writer}
+        self.scalar_dict = {}  # {writer_id : [[timestamp, step, value],...],...}
+
+    def __append_to_scalar_dict(self, tag, scalar_value, global_step,
+                                timestamp):
+        """This adds an entry to the self.scalar_dict datastructure with format
+        {writer_id : [[timestamp, step, value], ...], ...}.
+        """
+        from .x2num import makenp
+        if tag not in self.scalar_dict.keys():
+            self.scalar_dict[tag] = []
+        self.scalar_dict[tag].append([timestamp, global_step, float(makenp(scalar_value))])
+
+    def add_scalar(self, tag, scalar_value, global_step=None):
+        """Add scalar data to summary.
+
+        Args:
+            tag (string): Data identifier
+            scalar_value (float): Value to save
+            global_step (int): Global step value to record
+        """
+        self.file_writer.add_summary(scalar(tag, scalar_value), global_step)
+        self.__append_to_scalar_dict(tag, scalar_value, global_step, time.time())
+
+    def add_scalars(self, main_tag, tag_scalar_dict, global_step=None):
+        """Adds many scalar data to summary.
+
+        Args:
+            tag (string): Data identifier
+            main_tag (string): The parent name for the tags
+            tag_scalar_dict (dict): Key-value pair storing the tag and corresponding values
+            global_step (int): Global step value to record
+
+        Examples::
+
+            writer.add_scalars('run_14h',{'xsinx':i*np.sin(i/r),
+                                          'xcosx':i*np.cos(i/r),
+                                          'arctanx': numsteps*np.arctan(i/r)}, i)
+            # This function adds three values to the same scalar plot with the tag
+            # 'run_14h' in TensorBoard's scalar section.
+        """
+        timestamp = time.time()
+        fw_logdir = self.file_writer.get_logdir()
+        for tag, scalar_value in tag_scalar_dict.items():
+            fw_tag = fw_logdir + "/" + main_tag + "/" + tag
+            if fw_tag in self.all_writers.keys():
+                fw = self.all_writers[fw_tag]
+            else:
+                fw = FileWriter(logdir=fw_tag)
+                self.all_writers[fw_tag] = fw
+            fw.add_summary(scalar(main_tag, scalar_value), global_step)
+            self.__append_to_scalar_dict(fw_tag, scalar_value, global_step, timestamp)
+
+    def export_scalars_to_json(self, path):
+        """Exports to the given path an ASCII file containing all the scalars written
+        so far by this instance, with the following format:
+        {writer_id : [[timestamp, step, value], ...], ...}
+        """
+        with open(path, "w") as f:
+            json.dump(self.scalar_dict, f)
+
+    def add_histogram(self, tag, values, global_step=None, bins='tensorflow'):
+        """Add histogram to summary.
+
+        Args:
+            tag (string): Data identifier
+            values (numpy.array): Values to build histogram
+            global_step (int): Global step value to record
+            bins (string): one of {'tensorflow','auto', 'fd', ...}, this determines how the bins are made. You can find
+              other options in: https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html
+        """
+        if bins == 'tensorflow':
+            bins = self.default_bins
+        self.file_writer.add_summary(histogram(tag, values, bins), global_step)
+
+    def add_image(self, tag, img_tensor, global_step=None):
+        """Add image data to summary.
+
+        Note that this requires the ``pillow`` package.
+
+        Args:
+            tag (string): Data identifier
+            img_tensor (torch.Tensor): Image data
+            global_step (int): Global step value to record
+        Shape:
+            img_tensor: :math:`(3, H, W)`. Use ``torchvision.utils.make_grid()`` to prepare it is a good idea.
+        """
+        self.file_writer.add_summary(image(tag, img_tensor), global_step)
+
+    def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100):
+        """Add audio data to summary.
+
+        Args:
+            tag (string): Data identifier
+            snd_tensor (torch.Tensor): Sound data
+            global_step (int): Global step value to record
+            sample_rate (int): sample rate in Hz
+
+        Shape:
+            snd_tensor: :math:`(1, L)`. The values should lie between [-1, 1].
+        """
+        self.file_writer.add_summary(audio(tag, snd_tensor, sample_rate=sample_rate), global_step)
+
+    def add_text(self, tag, text_string, global_step=None):
+        """Add text data to summary.
+
+        Args:
+            tag (string): Data identifier
+            text_string (string): String to save
+            global_step (int): Global step value to record
+
+        Examples::
+
+            writer.add_text('lstm', 'This is an lstm', 0)
+            writer.add_text('rnn', 'This is an rnn', 10)
+        """
+        self.file_writer.add_summary(text(tag, text_string), global_step)
+        if tag not in self.text_tags:
+            self.text_tags.append(tag)
+            extension_dir = self.file_writer.get_logdir() + '/plugins/tensorboard_text/'
+            if not os.path.exists(extension_dir):
+                os.makedirs(extension_dir)
+            with open(extension_dir + 'tensors.json', 'w') as fp:
+                json.dump(self.text_tags, fp)
+
+    def add_graph_onnx(self, prototxt):
+        self.file_writer.add_graph_onnx(gg(prototxt))
+
+    def add_graph(self, model, input_to_model, verbose=False):
+        # prohibit second call?
+        # no, let tensorboard handles it and show its warning message.
+        """Add graph data to summary.
+
+        Args:
+            model (torch.nn.Module): model to draw.
+            input_to_model (torch.autograd.Variable): a variable or a tuple of variables to be fed.
+
+        """
+        import torch
+        from distutils.version import LooseVersion
+        if LooseVersion(torch.__version__) >= LooseVersion("0.3.1"):
+            pass
+        else:
+            if LooseVersion(torch.__version__) >= LooseVersion("0.3.0"):
+                print('You are using PyTorch==0.3.0, use add_graph_onnx()')
+                return
+            if not hasattr(torch.autograd.Variable, 'grad_fn'):
+                print('add_graph() only supports PyTorch v0.2.')
+                return
+        self.file_writer.add_graph(graph(model, input_to_model, verbose))
+
+    def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, tag='default'):
+        """Add embedding projector data to summary.
+
+        Args:
+            mat (torch.Tensor): A matrix which each row is the feature vector of the data point
+            metadata (list): A list of labels, each element will be convert to string
+            label_img (torch.Tensor): Images correspond to each data point
+            global_step (int): Global step value to record
+            tag (string): Name for the embedding
+        Shape:
+            mat: :math:`(N, D)`, where N is number of data and D is feature dimension
+
+            label_img: :math:`(N, C, H, W)`
+
+        Examples::
+
+            import keyword
+            import torch
+            meta = []
+            while len(meta)<100:
+                meta = meta+keyword.kwlist # get some strings
+            meta = meta[:100]
+
+            for i, v in enumerate(meta):
+                meta[i] = v+str(i)
+
+            label_img = torch.rand(100, 3, 10, 32)
+            for i in range(100):
+                label_img[i]*=i/100.0
+
+            writer.add_embedding(torch.randn(100, 5), metadata=meta, label_img=label_img)
+            writer.add_embedding(torch.randn(100, 5), label_img=label_img)
+            writer.add_embedding(torch.randn(100, 5), metadata=meta)
+        """
+        if global_step is None:
+            global_step = 0
+            # clear pbtxt?
+        save_path = os.path.join(self.file_writer.get_logdir(), str(global_step).zfill(5))
+        try:
+            os.makedirs(save_path)
+        except OSError:
+            print('warning: Embedding dir exists, did you set global_step for add_embedding()?')
+        if metadata is not None:
+            assert mat.size(0) == len(metadata), '#labels should equal with #data points'
+            make_tsv(metadata, save_path)
+        if label_img is not None:
+            assert mat.size(0) == label_img.size(0), '#images should equal with #data points'
+            make_sprite(label_img, save_path)
+        assert mat.dim() == 2, 'mat should be 2D, where mat.size(0) is the number of data points'
+        make_mat(mat.tolist(), save_path)
+        # new funcion to append to the config file a new embedding
+        append_pbtxt(metadata, label_img, self.file_writer.get_logdir(), str(global_step).zfill(5), tag)
+
+    def add_pr_curve(self, tag, labels, predictions, global_step=None, num_thresholds=127, weights=None):
+        """Adds precision recall curve.
+
+        Args:
+            tag (string): Data identifier
+            labels (torch.Tensor): Ground truth data. Binary label for each element.
+            predictions (torch.Tensor): The probability that an element be classified as true. Value should in [0, 1]
+            global_step (int): Global step value to record
+            num_thresholds (int): Number of thresholds used to draw the curve.
+
+        """
+        from .x2num import makenp
+        labels = makenp(labels)
+        predictions = makenp(predictions)
+        self.file_writer.add_summary(pr_curve(tag, labels, predictions, num_thresholds, weights), global_step)
+
+    def close(self):
+        if self.file_writer is None:
+            return  # ignore double close
+        self.file_writer.flush()
+        self.file_writer.close()
+        for path, writer in self.all_writers.items():
+            writer.flush()
+            writer.close()
+        self.file_writer = self.all_writers = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

tensorboardX/x2num.py

@@ -0,0 +1,90 @@
+# DO NOT alter/distruct/free input object !
+
+import numpy as np
+
+
+def makenp(x, modality=None):
+    # if already numpy, return
+    if isinstance(x, np.ndarray):
+        if modality == 'IMG' and x.dtype == np.uint8:
+            return x.astype(np.float32) / 255.0
+        return x
+    if np.isscalar(x):
+        return np.array([x])
+    if 'torch' in str(type(x)):
+        return pytorch_np(x, modality)
+    if 'chainer' in str(type(x)):
+        return chainer_np(x, modality)
+    if 'mxnet' in str(type(x)):
+        return mxnet_np(x, modality)
+
+
+def pytorch_np(x, modality):
+    import torch
+    if isinstance(x, torch.autograd.Variable):
+        x = x.data
+    x = x.cpu().numpy()
+    if modality == 'IMG':
+        x = _prepare_image(x)
+    return x
+
+
+def theano_np(x):
+    import theano
+    pass
+
+
+def caffe2_np(x):
+    pass
+
+
+def mxnet_np(x, modality):
+    x = x.asnumpy()
+    if modality == 'IMG':
+        x = _prepare_image(x)
+    return x
+
+
+def chainer_np(x, modality):
+    import chainer
+    x = chainer.cuda.to_cpu(x.data)
+    if modality == 'IMG':
+        x = _prepare_image(x)
+    return x
+
+
+def make_grid(I, ncols=8):
+    assert isinstance(I, np.ndarray), 'plugin error, should pass numpy array here'
+    assert I.ndim == 4 and I.shape[1] == 3
+    nimg = I.shape[0]
+    H = I.shape[2]
+    W = I.shape[3]
+    ncols = min(nimg, ncols)
+    nrows = int(np.ceil(float(nimg) / ncols))
+    canvas = np.zeros((3, H * nrows, W * ncols))
+    i = 0
+    for y in range(nrows):
+        for x in range(ncols):
+            if i >= nimg:
+                break
+            canvas[:, y * H:(y + 1) * H, x * W:(x + 1) * W] = I[i]
+            i = i + 1
+    return canvas
+
+
+def _prepare_image(I):
+    assert isinstance(I, np.ndarray), 'plugin error, should pass numpy array here'
+    assert I.ndim == 2 or I.ndim == 3 or I.ndim == 4
+    if I.ndim == 4:  # NCHW
+        if I.shape[1] == 1:  # N1HW
+            I = np.concatenate((I, I, I), 1)  # N3HW
+        assert I.shape[1] == 3
+        I = make_grid(I)  # 3xHxW
+    if I.ndim == 3 and I.shape[0] == 1:  # 1xHxW
+        I = np.concatenate((I, I, I), 0)  # 3xHxW
+    if I.ndim == 2:  # HxW
+        I = np.expand_dims(I, 0)  # 1xHxW
+        I = np.concatenate((I, I, I), 0)  # 3xHxW
+    I = I.transpose(1, 2, 0)
+
+    return I

test.py

@@ -15,8 +15,9 @@ from PIL import Image
 import lib.TestTransforms as transforms
 from models.ColorVidNet import ColorVidNet
 from models.FrameColor import frame_colorization
-from models.NonlocalNet import VGG19_feature_color_torchversion, WarpNet
-from utils.util import batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not, save_frames, tensor_lab2rgb, uncenter_l
+from models.NonlocalNet import VGG19_pytorch, WarpNet
+from utils.util import (batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not,
+                        save_frames, tensor_lab2rgb, uncenter_l)
 from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor
 
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
@@ -36,7 +37,7 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c
     print("processing the folder:", input_path)
     path, dirs, filenames = os.walk(input_path).__next__()
     file_count = len(filenames)
-    filenames.sort()
+    filenames.sort(key=lambda f: int("".join(filter(str.isdigit, f) or -1)))
 
     # NOTE: resize frames to 216*384
     transform = transforms.Compose(
@@ -129,6 +130,7 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c
 
         # save the frames
         save_frames(IA_predict_rgb, output_path, index)
+
     # output video
     video_name = "video.avi"
     folder2vid(image_folder=output_path, output_dir=output_path, filename=video_name)
@@ -145,7 +147,7 @@ if __name__ == "__main__":
     parser.add_argument("--clip_path", type=str, default="./sample_videos/clips/v32", help="path of input clips")
     parser.add_argument("--ref_path", type=str, default="./sample_videos/ref/v32", help="path of refernce images")
     parser.add_argument("--output_path", type=str, default="./sample_videos/output", help="path of output clips")
-    opt = parser.parse_args(args=[])
+    opt = parser.parse_args()
     opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")]
     cudnn.benchmark = True
     print("running on GPU", opt.gpu_ids)
@@ -156,7 +158,7 @@ if __name__ == "__main__":
 
     nonlocal_net = WarpNet(1)
     colornet = ColorVidNet(7)
-    vggnet = VGG19_feature_color_torchversion()
+    vggnet = VGG19_pytorch()
     vggnet.load_state_dict(torch.load("data/vgg19_conv.pth"))
     for param in vggnet.parameters():
         param.requires_grad = False

train.py

@@ -0,0 +1,738 @@
+from __future__ import print_function
+
+import argparse
+import math
+import os
+import queue
+import time
+
+import cv2
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transform_lib
+from torch.utils.data import ConcatDataset, DataLoader, WeightedRandomSampler
+from torchvision.transforms import CenterCrop
+
+import lib.TrainTransforms as transforms
+from lib.videoloader import VideosDataset
+from lib.videoloader_imagenet import VideosDataset_ImageNet
+from models.ColorVidNet import ColorVidNet
+from models.ContextualLoss import ContextualLoss, ContextualLoss_forward
+from models.FrameColor import frame_colorization
+from models.GAN_models import Discriminator_x64
+from models.NonlocalNet import (NonlocalWeightedAverage, VGG19_pytorch,
+                                WarpNet, WeightedAverage,
+                                WeightedAverage_color)
+from tensorboardX import SummaryWriter
+from utils.util import (batch_lab2rgb_transpose_mc, feature_normalize, l1_loss,
+                        mkdir_if_not, mse_loss, parse, tensor_lab2rgb,
+                        uncenter_l, weighted_l1_loss, weighted_mse_loss)
+from utils.util_distortion import (CenterPad_threshold, Normalize, RGB2Lab,
+                                   ToTensor)
+from utils.util_tensorboard import TBImageRecorder, value_logger
+from utils.warping import WarpingLayer
+
+cv2.setNumThreads(0)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--data_root", default="/home/zhanbo/remote/video/video_pair3/", type=str)
+parser.add_argument("--data_root_imagenet", default="/home/zhanbo/remote/imagenet_ref_pair/", type=str)
+parser.add_argument("--gpu_ids", type=str, default="0,1,2,3", help="separate by comma")
+parser.add_argument("--workers", type=int, default=8)
+parser.add_argument("--batch_size", type=int, default=16)
+parser.add_argument("--image_size", type=int, default=[216, 384])
+parser.add_argument("--ic", type=int, default=7)
+parser.add_argument("--epoch", type=int, default=40)
+
+parser.add_argument("--resume_epoch", type=int, default=0)
+parser.add_argument("--resume", type=bool, default=False)
+parser.add_argument("--load_pretrained_model", type=bool, default=True)
+
+parser.add_argument("--lr", type=float, default=0.00001)
+parser.add_argument("--beta1", type=float, default=0.5)
+parser.add_argument("--lr_step", type=int, default=100)
+parser.add_argument("--lr_gamma", type=float, default=0.1)
+
+
+parser.add_argument("--checkpoint_dir", type=str, default="checkpoints/default")
+parser.add_argument("--tb_log_step", type=int, default=50)
+parser.add_argument("--print_step", type=int, default=2)
+
+parser.add_argument("--real_reference_probability", type=float, default=0.7)
+parser.add_argument("--nonzero_placeholder_probability", type=float, default=0.0)
+parser.add_argument("--with_bad", type=bool, default=True)
+parser.add_argument("--with_mid", type=bool, default=True)
+
+parser.add_argument("--domain_invariant", type=bool, default=False)
+parser.add_argument("--weigth_l1", type=float, default=2.0)
+parser.add_argument("--weight_contextual", type=float, default="0.2")
+parser.add_argument("--weight_perceptual", type=float, default="0.001")
+parser.add_argument("--weight_smoothness", type=float, default="5.0")
+parser.add_argument("--weight_gan", type=float, default="0.2")
+parser.add_argument("--weight_nonlocal_smoothness", type=float, default="0.0")
+parser.add_argument("--weight_nonlocal_consistent", type=float, default="0.0")
+parser.add_argument("--weight_consistent", type=float, default="0.02")
+parser.add_argument("--luminance_noise", type=float, default="2.0")
+parser.add_argument("--permute_data", type=bool, default=True)
+parser.add_argument("--contextual_loss_direction", type=str, default="forward", help="forward or backward matching")
+
+
+def image_logger_fn(
+    I_last_lab,
+    I_current_lab,
+    I_reference_lab,
+    I_last_lab_predict,
+    I_current_lab_predict,
+    I_last_nonlocal_lab,
+    I_current_nonlocal_lab,
+):
+    I_last_image = batch_lab2rgb_transpose_mc(I_last_lab[0:32, 0:1, :, :], I_last_lab[0:32, 1:3, :, :])
+    I_current_image = batch_lab2rgb_transpose_mc(I_current_lab[0:32, 0:1, :, :], I_current_lab[0:32, 1:3, :, :])
+    I_reference_image = batch_lab2rgb_transpose_mc(I_reference_lab[0:32, 0:1, :, :], I_reference_lab[0:32, 1:3, :, :])
+    I_last_image_predict = batch_lab2rgb_transpose_mc(
+        I_last_lab_predict[0:32, 0:1, :, :], I_last_lab_predict[0:32, 1:3, :, :]
+    )
+    I_current_image_predict = batch_lab2rgb_transpose_mc(
+        I_current_lab_predict[0:32, 0:1, :, :], I_current_lab_predict[0:32, 1:3, :, :]
+    )
+    I_last_nonlocal_image = batch_lab2rgb_transpose_mc(
+        I_last_nonlocal_lab[0:32, 0:1, :, :], I_last_nonlocal_lab[0:32, 1:3, :, :]
+    )
+    I_current_nonlocal_image = batch_lab2rgb_transpose_mc(
+        I_current_nonlocal_lab[0:32, 0:1, :, :], I_current_nonlocal_lab[0:32, 1:3, :, :]
+    )
+
+    img_info = {}
+    img_info["1_I_last"] = I_last_image
+    img_info["2_I_current"] = I_current_image
+    img_info["3_I_reference"] = I_reference_image
+    img_info["4_I_last_predict"] = I_last_image_predict
+    img_info["5_I_curren_predict"] = I_current_image_predict
+    img_info["6_I_last_nonlocal"] = I_last_nonlocal_image
+    img_info["7_I_current_nonlocal"] = I_current_nonlocal_image
+
+    return img_info
+
+
+def training_logger():
+    try:
+        if total_iter % opt.print_step == 0:
+            print("processing time:", elapsed)
+            print(
+                "Epoch %d, Step[%d/%d], lr: %f, total_loss: %.2f"
+                % (
+                    epoch,
+                    ((iter + 1) % iter_num_per_epoch),
+                    iter_num_per_epoch,
+                    step_optim_scheduler_g.get_last_lr()[0],
+                    total_loss.item(),
+                )
+            )
+
+            value_logger(
+                tb_writer,
+                total_iter,
+                loss_info={
+                    "l1_loss": l1_loss.item(),
+                    "feat_loss": feat_loss.item(),
+                    "contextual_loss_total": contextual_loss_total.item(),
+                    "smoothness_loss": smoothness_loss.item(),
+                    "nonlocal_smoothness_loss": nonlocal_smoothness_loss.item(),
+                    "nonlocal_consistent_loss": nonlocal_consistent_loss.item(),
+                    "consistent_loss": consistent_loss.item(),
+                    "generator_loss": generator_loss.item(),
+                    "discriminator_loss": discriminator_loss.item(),
+                    "total_loss": total_loss.item(),
+                },
+            )
+
+        if total_iter % opt.tb_log_step == 0:
+            I_last_nonlocal_lab = torch.cat(
+                (I_last_lab[:, 0:1, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :]), dim=1
+            )
+            I_current_nonlocal_lab = torch.cat(
+                (I_current_lab[:, 0:1, :, :], I_current_nonlocal_lab_predict[:, 1:3, :, :]), dim=1
+            )
+            I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1)
+            data_queue.put(
+                (
+                    (
+                        I_last_lab.cpu(),
+                        I_current_lab.cpu(),
+                        I_reference_lab.cpu(),
+                        I_last_lab_predict.cpu(),
+                        I_current_lab_predict.cpu(),
+                        I_last_nonlocal_lab.cpu(),
+                        I_current_nonlocal_lab,
+                    ),
+                    total_iter,
+                )
+            )
+
+        if total_iter % 2000 == 0:
+            if len(opt.gpu_ids) > 1:
+                torch.save(
+                    nonlocal_net.module.state_dict(),
+                    os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter,
+                )
+                torch.save(
+                    colornet.module.state_dict(),
+                    os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter,
+                )
+                torch.save(
+                    discriminator.module.state_dict(),
+                    os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter,
+                )
+            else:
+                torch.save(
+                    nonlocal_net.state_dict(),
+                    os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter,
+                )
+                torch.save(colornet.state_dict(), os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter)
+                torch.save(
+                    discriminator.state_dict(),
+                    os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter,
+                )
+
+        # save the state for resume
+        if total_iter % 2000 == 0:
+            print("saving the checkpoint")
+            if len(opt.gpu_ids) > 1:
+                state = {
+                    "total_iter": total_iter + 1,
+                    "epoch": epoch,
+                    "colornet_state": colornet.module.state_dict(),
+                    "nonlocal_net_state": nonlocal_net.module.state_dict(),
+                    "discriminator_state": discriminator.module.state_dict(),
+                    "optimizer_g": optimizer_g.state_dict(),
+                    "optimizer_d": optimizer_d.state_dict(),
+                    "optimizer_schedule_g": step_optim_scheduler_g.state_dict(),
+                    "optimizer_schedule_d": step_optim_scheduler_g.state_dict(),
+                }
+            else:
+                state = {
+                    "total_iter": total_iter + 1,
+                    "epoch": epoch,
+                    "colornet_state": colornet.state_dict(),
+                    "nonlocal_net_state": nonlocal_net.state_dict(),
+                    "discriminator_state": discriminator.state_dict(),
+                    "optimizer_g": optimizer_g.state_dict(),
+                    "optimizer_d": optimizer_d.state_dict(),
+                    "optimizer_schedule_g": step_optim_scheduler_g.state_dict(),
+                    "optimizer_schedule_d": step_optim_scheduler_d.state_dict(),
+                }
+            torch.save(state, os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth"))
+    except Exception as e:
+        print("Exception during output")
+        print(e)
+
+
+def gpu_setup():
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+    cudnn.benchmark = True
+    opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")]
+    torch.cuda.set_device(opt.gpu_ids[0])
+    device = torch.device("cuda")
+    print("running on GPU", opt.gpu_ids)
+    return device
+
+
+def load_data():
+    print("initializing dataloader")
+    transforms_video = [
+        CenterCrop(opt.image_size),
+        RGB2Lab(),
+        ToTensor(),
+        Normalize(),
+    ]
+    transforms_imagenet = [CenterPad_threshold(opt.image_size), RGB2Lab(), ToTensor(), Normalize()]
+    extra_reference_transform = [
+        transform_lib.RandomHorizontalFlip(0.5),
+        transform_lib.RandomResizedCrop(480, (0.98, 1.0), ratio=(0.8, 1.2)),
+    ]
+    train_dataset_video = VideosDataset(
+        data_root=opt.data_root,
+        epoch=opt.epoch,
+        image_size=opt.image_size,
+        image_transform=transforms.Compose(transforms_video),
+        real_reference_probability=opt.real_reference_probability,
+        nonzero_placeholder_probability=opt.nonzero_placeholder_probability,
+    )
+    train_dataset_imagenet = VideosDataset_ImageNet(
+        data_root=opt.data_root_imagenet,
+        image_size=opt.image_size,
+        epoch=opt.epoch,
+        with_bad=opt.with_bad,
+        with_mid=opt.with_mid,
+        transforms_imagenet=transforms_imagenet,
+        distortion_level=4,
+        brightnessjitter=5,
+        nonzero_placeholder_probability=opt.nonzero_placeholder_probability,
+        extra_reference_transform=extra_reference_transform,
+        real_reference_probability=opt.real_reference_probability,
+    )
+
+    video_training_length = len(train_dataset_video)
+    imagenet_training_length = len(train_dataset_imagenet)
+    dataset_training_length = train_dataset_video.real_len + train_dataset_imagenet.real_len
+    dataset_combined = ConcatDataset([train_dataset_video, train_dataset_imagenet])
+    sampler = WeightedRandomSampler(
+        [1] * video_training_length + [1] * imagenet_training_length, dataset_training_length * opt.epoch
+    )
+    data_loader = DataLoader(
+        dataset_combined,
+        batch_size=opt.batch_size,
+        shuffle=False,
+        num_workers=opt.workers,
+        pin_memory=True,
+        drop_last=True,
+        sampler=sampler,
+    )
+    return dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader
+
+
+def define_loss():
+    print("defining loss")
+    ab_criterion = nn.SmoothL1Loss().to(device)
+    nonlocal_criterion = nn.SmoothL1Loss().to(device)
+    feat_l2_criterion = nn.MSELoss().to(device)
+    feat_l1_criterion = nn.SmoothL1Loss().to(device)
+    contextual_loss = ContextualLoss().to(device)
+    contextual_forward_loss = ContextualLoss_forward().to(device)
+    BCE_stable = nn.BCEWithLogitsLoss().to(device)
+    return contextual_loss, contextual_forward_loss
+
+
+def define_optimizer():
+    print("defining optimizer")
+    optimizer_g = optim.Adam(
+        [{"params": nonlocal_net.parameters(), "lr": 1e-5}, {"params": colornet.parameters(), "lr": 2e-4}],
+        betas=(0.5, 0.999),
+        eps=1e-5,
+        amsgrad=True,
+    )
+    optimizer_d = optim.Adam(
+        filter(lambda p: p.requires_grad, discriminator.parameters()), lr=2 * 1e-4, betas=(0.5, 0.999)
+    )
+    return optimizer_g, optimizer_d
+
+
+def resume_model():
+    print("resuming the learning")
+    checkpoint = torch.load(os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth"))
+    total_iter = checkpoint["total_iter"]
+    epoch = checkpoint["epoch"]
+    colornet.load_state_dict(checkpoint["colornet_state"])
+    nonlocal_net.load_state_dict(checkpoint["nonlocal_net_state"])
+    discriminator.load_state_dict(checkpoint["discriminator_state"])
+    optimizer_g.load_state_dict(checkpoint["optimizer_g"])
+    optimizer_d.load_state_dict(checkpoint["optimizer_d"])
+    step_optim_scheduler_g.load_state_dict(checkpoint["optimizer_schedule_g"])
+    step_optim_scheduler_d.load_state_dict(checkpoint["optimizer_schedule_d"])
+
+
+def to_device(
+    colornet,
+    nonlocal_net,
+    discriminator,
+    vggnet,
+    contextual_loss,
+    contextual_forward_loss,
+    weighted_layer_color,
+    nonlocal_weighted_layer,
+    warping_layer,
+    instancenorm,
+):
+    print("moving models to device")
+    colornet = torch.nn.DataParallel(colornet.to(device), device_ids=opt.gpu_ids)
+    nonlocal_net = torch.nn.DataParallel(nonlocal_net.to(device), device_ids=opt.gpu_ids)
+    discriminator = torch.nn.DataParallel(discriminator.to(device), device_ids=opt.gpu_ids)
+    vggnet = torch.nn.DataParallel(vggnet.to(device), device_ids=opt.gpu_ids)
+    contextual_loss = torch.nn.DataParallel(contextual_loss.to(device), device_ids=opt.gpu_ids)
+    contextual_forward_loss = torch.nn.DataParallel(contextual_forward_loss.to(device), device_ids=opt.gpu_ids)
+    weighted_layer_color = torch.nn.DataParallel(weighted_layer_color.to(device), device_ids=opt.gpu_ids)
+    nonlocal_weighted_layer = torch.nn.DataParallel(nonlocal_weighted_layer.to(device), device_ids=opt.gpu_ids)
+    warping_layer = torch.nn.DataParallel(warping_layer.to(device), device_ids=opt.gpu_ids)
+    instancenorm = torch.nn.DataParallel(instancenorm.to(device), device_ids=opt.gpu_ids)
+
+    return (
+        vggnet,
+        nonlocal_net,
+        colornet,
+        discriminator,
+        instancenorm,
+        contextual_loss,
+        contextual_forward_loss,
+        weighted_layer_color,
+        nonlocal_weighted_layer,
+        warping_layer,
+    )
+
+
+def loss_init():
+    print("initializing losses")
+    zero_loss = torch.Tensor([0]).to(device)
+    (
+        feat_loss,
+        contextual_loss_total,
+        smoothness_loss,
+        nonlocal_smoothness_loss,
+        consistent_loss,
+        nonlocal_consistent_loss,
+        generator_loss,
+        discriminator_loss,
+    ) = (zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss)
+
+    return (
+        feat_loss,
+        contextual_loss_total,
+        smoothness_loss,
+        nonlocal_smoothness_loss,
+        consistent_loss,
+        nonlocal_consistent_loss,
+        generator_loss,
+        discriminator_loss,
+    )
+
+
+def video_colorization():
+    # colorization for the last frame
+    I_last_ab_predict, I_last_nonlocal_lab_predict, features_last_gray = frame_colorization(
+        I_last_lab,
+        I_reference_lab,
+        placeholder_lab,
+        features_B,
+        vggnet,
+        nonlocal_net,
+        colornet,
+        feature_noise=0,
+        luminance_noise=opt.luminance_noise,
+    )
+    I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1)
+
+    # colorization for the current frame
+    I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization(
+        I_current_lab,
+        I_reference_lab,
+        I_last_lab_predict,
+        features_B,
+        vggnet,
+        nonlocal_net,
+        colornet,
+        feature_noise=0,
+        luminance_noise=opt.luminance_noise,
+    )
+    I_current_lab_predict = torch.cat((I_last_l, I_current_ab_predict), dim=1)
+    return I_current_ab_predict, I_last_ab_predict, I_current_nonlocal_lab_predict, I_last_nonlocal_lab_predict
+
+
+if __name__ == "__main__":
+    torch.multiprocessing.set_start_method("spawn", force=True)
+    opt = parse(parser)
+    opt.data_root = opt.data_root.split(",")[0]
+    opt.data_root_imagenet = opt.data_root_imagenet.split(",")[0]
+    mkdir_if_not(opt.checkpoint_dir)
+    mkdir_if_not("./runs/")
+
+    device = gpu_setup()
+    dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader = load_data()
+    tb_writer = SummaryWriter()
+    data_queue = queue.Queue()
+    tb_image_reorder = TBImageRecorder(tb_writer, image_logger_fn, data_queue)
+    tb_image_reorder.start()
+
+    # define network
+    nonlocal_net = WarpNet(opt.batch_size)
+    colornet = ColorVidNet(opt.ic)
+    discriminator = Discriminator_x64(ndf=64)
+    colornet.to(device)
+    nonlocal_net.to(device)
+    discriminator.to(device)
+
+    weighted_layer = WeightedAverage()
+    weighted_layer_color = WeightedAverage_color()
+    nonlocal_weighted_layer = NonlocalWeightedAverage()
+    instancenorm = nn.InstanceNorm2d(512, affine=False)
+    warping_layer = WarpingLayer("gpu")
+
+    vggnet = VGG19_pytorch()
+    vggnet.load_state_dict(torch.load("data/vgg19_conv.pth"))
+    vggnet.eval()
+    for param in vggnet.parameters():
+        param.requires_grad = False
+
+    # load pre-trained model
+    if opt.load_pretrained_model:
+        nonlocal_pretain_path = os.path.join("checkpoints/video_moredata_l1/", "nonlocal_net_iter_76000.pth")
+        nonlocal_net.load_state_dict(torch.load(nonlocal_pretain_path))
+        color_test_path = "checkpoints/video_moredata_l1/" + "colornet_iter_76000.pth"
+        colornet.load_state_dict(torch.load(color_test_path))
+
+    # define loss function
+    contextual_loss, contextual_forward_loss = define_loss()
+
+    # define optimizer
+    optimizer_g, optimizer_d = define_optimizer()
+    step_optim_scheduler_g = optim.lr_scheduler.StepLR(optimizer_g, step_size=opt.lr_step, gamma=opt.lr_gamma)
+    step_optim_scheduler_d = optim.lr_scheduler.StepLR(optimizer_d, step_size=opt.lr_step, gamma=opt.lr_gamma)
+
+    # define others
+    downsampling_by2 = nn.AvgPool2d(kernel_size=2).to(device)
+    downsampling_by4 = nn.AvgPool2d(kernel_size=4).to(device)
+
+    # dataset info
+    iter_num_per_epoch = dataset_training_length // opt.batch_size
+    total_iter = opt.resume_epoch * iter_num_per_epoch
+    print(
+        "train_dataset info,  real_len: %d, epoch_len: %d, iter_num_per_epoch: %d"
+        % (dataset_training_length, len(train_dataset_video) + len(train_dataset_imagenet), iter_num_per_epoch)
+    )
+
+    if opt.resume:
+        resume_model()
+
+    # move to GPU processing
+    device = opt.gpu_ids[0]
+    (
+        vggnet,
+        nonlocal_net,
+        colornet,
+        discriminator,
+        instancenorm,
+        contextual_loss,
+        contextual_forward_loss,
+        weighted_layer_color,
+        nonlocal_weighted_layer,
+        warping_layer,
+    ) = to_device(
+        colornet,
+        nonlocal_net,
+        discriminator,
+        vggnet,
+        contextual_loss,
+        contextual_forward_loss,
+        weighted_layer_color,
+        nonlocal_weighted_layer,
+        warping_layer,
+        instancenorm,
+    )
+
+    (
+        feat_loss,
+        contextual_loss_total,
+        smoothness_loss,
+        nonlocal_smoothness_loss,
+        consistent_loss,
+        nonlocal_consistent_loss,
+        generator_loss,
+        discriminator_loss,
+    ) = loss_init()
+
+    # %% Training
+    print("start training")
+    for iter, data in enumerate(data_loader):
+        start_time = time.time()
+        total_iter += 1
+        epoch = math.ceil(total_iter / iter_num_per_epoch)
+
+        ###### LOADING DATA SAMPLE ######
+        (
+            I_last_lab,
+            I_current_lab,
+            I_reference_lab,
+            flow_forward,
+            flow_backward,
+            mask,
+            placeholder_lab,
+            self_ref_flag,
+        ) = data
+        I_last_lab = I_last_lab.to(device)
+        I_current_lab = I_current_lab.to(device)
+        I_reference_lab = I_reference_lab.to(device)
+        flow_forward = flow_forward.to(device)
+        flow_backward = flow_backward.to(device)
+        mask = mask.to(device)
+        placeholder_lab = placeholder_lab.to(device)
+        self_ref_flag = self_ref_flag.to(device)
+
+        I_last_l = I_last_lab[:, 0:1, :, :]
+        I_last_ab = I_last_lab[:, 1:3, :, :]
+        I_current_l = I_current_lab[:, 0:1, :, :]
+        I_current_ab = I_current_lab[:, 1:3, :, :]
+        I_reference_l = I_reference_lab[:, 0:1, :, :]
+        I_reference_ab = I_reference_lab[:, 1:3, :, :]
+        I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1))
+        features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True)
+
+        ###### COLORIZATION ######
+        (
+            I_current_ab_predict,
+            I_last_ab_predict,
+            I_current_nonlocal_lab_predict,
+            I_last_nonlocal_lab_predict,
+        ) = video_colorization()
+
+        ###### UPDATE DISCRIMINATOR ######
+        if opt.weight_gan > 0:
+            optimizer_g.zero_grad()
+            optimizer_d.zero_grad()
+
+            fake_data_lab = torch.cat(
+                (uncenter_l(I_current_l), I_current_ab_predict, uncenter_l(I_last_l), I_last_ab_predict), dim=1
+            )
+            real_data_lab = torch.cat((uncenter_l(I_current_l), I_current_ab, uncenter_l(I_last_l), I_last_ab), dim=1)
+
+            if opt.permute_data:
+                batch_index = torch.arange(-1, opt.batch_size - 1, dtype=torch.long)
+                real_data_lab = real_data_lab[batch_index, ...]
+
+            y_pred_fake, feature_pred_fake = discriminator(fake_data_lab.detach())
+            y_pred_real, feature_pred_real = discriminator(real_data_lab.detach())
+
+            y = torch.ones_like(y_pred_real)
+            y2 = torch.zeros_like(y_pred_real)
+            discriminator_loss = (
+                torch.mean((y_pred_real - torch.mean(y_pred_fake) - y) ** 2)
+                + torch.mean((y_pred_fake - torch.mean(y_pred_real) + y) ** 2)
+            ) / 2
+            discriminator_loss.backward()
+            optimizer_d.step()
+
+        ###### UPDATE GENERATOR ######
+        optimizer_g.zero_grad()
+        optimizer_d.zero_grad()
+
+        # extract vgg features for both output and original image
+        I_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab_predict), dim=1))
+        predict_relu1_1, predict_relu2_1, predict_relu3_1, predict_relu4_1, predict_relu5_1 = vggnet(
+            I_predict_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True
+        )
+
+        I_current_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab), dim=1))
+        A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet(
+            I_current_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True
+        )
+        B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B
+
+        ###### LOSS COMPUTE ######
+        # l1 loss
+        if opt.weigth_l1 > 0:
+            sample_weights = (self_ref_flag[:, 1:3, :, :]) / (sum(self_ref_flag[:, 0, 0, 0]) + 1e-5)
+            l1_loss = weighted_l1_loss(I_current_ab_predict, I_current_ab, sample_weights) * opt.weigth_l1
+
+        # generator loss
+        if opt.weight_gan > 0:
+            y_pred_fake, feature_pred_fake = discriminator(fake_data_lab)
+            y_pred_real, feature_pred_real = discriminator(real_data_lab)
+            generator_loss = (
+                (
+                    torch.mean((y_pred_real - torch.mean(y_pred_fake) + y) ** 2)
+                    + torch.mean((y_pred_fake - torch.mean(y_pred_real) - y) ** 2)
+                )
+                / 2
+                * opt.weight_gan
+            )
+
+        # feature loss
+        if opt.domain_invariant:
+            feat_loss = (
+                mse_loss(instancenorm(predict_relu5_1), instancenorm(A_relu5_1.detach()))
+                * opt.weight_perceptual
+                * 1e5
+                * 0.2
+            )
+        else:
+            feat_loss = mse_loss(predict_relu5_1, A_relu5_1.detach()) * opt.weight_perceptual
+
+        # contextual loss
+        if opt.contextual_loss_direction == "backward":
+            contextual_style5_1 = torch.mean(contextual_loss(predict_relu5_1, B_relu5_1.detach())) * 8
+            contextual_style4_1 = torch.mean(contextual_loss(predict_relu4_1, B_relu4_1.detach())) * 4
+            contextual_style3_1 = (
+                torch.mean(contextual_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))) * 2
+            )
+        else:
+            contextual_style5_1 = torch.mean(contextual_forward_loss(predict_relu5_1, B_relu5_1.detach())) * 8
+            contextual_style4_1 = torch.mean(contextual_forward_loss(predict_relu4_1, B_relu4_1.detach())) * 4
+            contextual_style3_1 = (
+                torch.mean(
+                    contextual_forward_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))
+                )
+                * 2
+            )
+        if opt.weight_contextual > 0:
+            contextual_loss_total = (
+                contextual_style5_1 + contextual_style4_1 + contextual_style3_1
+            ) * opt.weight_contextual
+
+        # smoothness loss
+        if opt.weight_smoothness > 0:
+            scale_factor = 1
+            I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1)
+            IA_ab_weighed = weighted_layer_color(
+                I_current_lab, I_current_lab_predict, patch_size=3, alpha=10, scale_factor=scale_factor
+            )
+            smoothness_loss = (
+                mse_loss(nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), IA_ab_weighed)
+                * opt.weight_smoothness
+            )
+
+        if opt.weight_nonlocal_smoothness > 0:
+            scale_factor = 0.25
+            alpha_nonlocal_smoothness = 0.5
+            nonlocal_smooth_feature = feature_normalize(A_relu2_1)
+            I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1)
+            I_current_ab_weighted_nonlocal = nonlocal_weighted_layer(
+                I_current_lab_predict,
+                nonlocal_smooth_feature.detach(),
+                patch_size=3,
+                alpha=alpha_nonlocal_smoothness,
+                scale_factor=scale_factor,
+            )
+            nonlocal_smoothness_loss = (
+                mse_loss(
+                    nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor),
+                    I_current_ab_weighted_nonlocal,
+                )
+                * opt.weight_nonlocal_smoothness
+            )
+
+        if opt.weight_consistent:
+            I_current_lab_predict_warp = warping_layer(I_current_lab_predict, flow_forward)
+            I_current_ab_predict_warp = I_current_lab_predict_warp[:, 1:3, :, :]
+            consistent_loss = (
+                weighted_mse_loss(I_current_ab_predict_warp, I_last_ab_predict, mask) * opt.weight_consistent
+            )
+
+        if opt.weight_nonlocal_consistent:
+            I_current_nonlocal_lab_predict_warp = warping_layer(I_current_nonlocal_lab_predict, flow_forward)
+            nonlocal_consistent_loss = (
+                weighted_mse_loss(
+                    I_current_nonlocal_lab_predict_warp[:, 1:3, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :], mask
+                )
+                * opt.weight_nonlocal_consistent
+            )
+
+        # total loss
+        total_loss = (
+            l1_loss
+            + feat_loss
+            + contextual_loss_total
+            + smoothness_loss
+            + nonlocal_smoothness_loss
+            + consistent_loss
+            + nonlocal_consistent_loss
+            + generator_loss
+        )
+        total_loss.backward()
+        optimizer_g.step()
+        end_time = time.time()
+        elapsed = end_time - start_time
+        training_logger()
+
+        step_optim_scheduler_g.step()
+        step_optim_scheduler_d.step()
+
+    data_queue.put((None, None))

utils/flowlib.py

@@ -0,0 +1,523 @@
+#!/usr/bin/python
+"""
+# ==============================
+# flowlib.py
+# library for optical flow processing
+# Author: Ruoteng Li
+# Date: 6th Aug 2016
+# ==============================
+"""
+import matplotlib.colors as cl
+import matplotlib.pyplot as plt
+import numpy as np
+import png
+from PIL import Image
+
+UNKNOWN_FLOW_THRESH = 1e7
+SMALLFLOW = 0.0
+LARGEFLOW = 1e8
+
+"""
+=============
+Flow Section
+=============
+"""
+
+
+def show_flow(filename):
+    """
+    visualize optical flow map using matplotlib
+    :param filename: optical flow file
+    :return: None
+    """
+    flow = read_flow(filename)
+    img = flow_to_image(flow)
+    plt.imshow(img)
+    plt.show()
+
+
+def visualize_flow(flow, mode="Y"):
+    """
+    this function visualize the input flow
+    :param flow: input flow in array
+    :param mode: choose which color mode to visualize the flow (Y: Ccbcr, RGB: RGB color)
+    :return: None
+    """
+    if mode == "Y":
+        # Ccbcr color wheel
+        img = flow_to_image(flow)
+        plt.imshow(img)
+        plt.show()
+    elif mode == "RGB":
+        (h, w) = flow.shape[0:2]
+        du = flow[:, :, 0]
+        dv = flow[:, :, 1]
+        valid = flow[:, :, 2]
+        max_flow = max(np.max(du), np.max(dv))
+        img = np.zeros((h, w, 3), dtype=np.float64)
+        # angle layer
+        img[:, :, 0] = np.arctan2(dv, du) / (2 * np.pi)
+        # magnitude layer, normalized to 1
+        img[:, :, 1] = np.sqrt(du * du + dv * dv) * 8 / max_flow
+        # phase layer
+        img[:, :, 2] = 8 - img[:, :, 1]
+        # clip to [0,1]
+        small_idx = img[:, :, 0:3] < 0
+        large_idx = img[:, :, 0:3] > 1
+        img[small_idx] = 0
+        img[large_idx] = 1
+        # convert to rgb
+        img = cl.hsv_to_rgb(img)
+        # remove invalid point
+        img[:, :, 0] = img[:, :, 0] * valid
+        img[:, :, 1] = img[:, :, 1] * valid
+        img[:, :, 2] = img[:, :, 2] * valid
+        # show
+        plt.imshow(img)
+        plt.show()
+
+    return None
+
+
+def read_flow(filename):
+    """
+    read optical flow from Middlebury .flo file
+    :param filename: name of the flow file
+    :return: optical flow data in matrix
+    """
+    f = open(filename, "rb")
+    try:
+        magic = np.fromfile(f, np.float32, count=1)[0]  # For Python3.x
+    except:
+        magic = np.fromfile(f, np.float32, count=1)  # For Python2.x
+    data2d = None
+
+    if 202021.25 != magic:
+        print("Magic number incorrect. Invalid .flo file")
+    else:
+        w = np.fromfile(f, np.int32, count=1)[0]
+        h = np.fromfile(f, np.int32, count=1)[0]
+        # print("Reading %d x %d flo file" % (h, w))
+        data2d = np.fromfile(f, np.float32, count=2 * w * h)
+        # reshape data into 3D array (columns, rows, channels)
+        data2d = np.resize(data2d, (h, w, 2))
+    f.close()
+    return data2d
+
+
+def read_flow_png(flow_file):
+    """
+    Read optical flow from KITTI .png file
+    :param flow_file: name of the flow file
+    :return: optical flow data in matrix
+    """
+    flow_object = png.Reader(filename=flow_file)
+    flow_direct = flow_object.asDirect()
+    flow_data = list(flow_direct[2])
+    (w, h) = flow_direct[3]["size"]
+    flow = np.zeros((h, w, 3), dtype=np.float64)
+    for i in range(len(flow_data)):
+        flow[i, :, 0] = flow_data[i][0::3]
+        flow[i, :, 1] = flow_data[i][1::3]
+        flow[i, :, 2] = flow_data[i][2::3]
+
+    invalid_idx = flow[:, :, 2] == 0
+    flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0
+    flow[invalid_idx, 0] = 0
+    flow[invalid_idx, 1] = 0
+    return flow
+
+
+def write_flow(flow, filename):
+    """
+    write optical flow in Middlebury .flo format
+    :param flow: optical flow map
+    :param filename: optical flow file path to be saved
+    :return: None
+    """
+    with open(filename, "wb") as f:
+        magic = np.array([202021.25], dtype=np.float32)
+        (height, width) = flow.shape[0:2]
+        w = np.array([width], dtype=np.int32)
+        h = np.array([height], dtype=np.int32)
+        magic.tofile(f)
+        w.tofile(f)
+        h.tofile(f)
+        flow.tofile(f)
+
+
+def segment_flow(flow):
+    h = flow.shape[0]
+    w = flow.shape[1]
+    u = flow[:, :, 0]
+    v = flow[:, :, 1]
+
+    idx = (abs(u) > LARGEFLOW) | (abs(v) > LARGEFLOW)
+    idx2 = abs(u) == SMALLFLOW
+    class0 = (v == 0) & (u == 0)
+    u[idx2] = 0.00001
+    tan_value = v / u
+
+    class1 = (tan_value < 1) & (tan_value >= 0) & (u > 0) & (v >= 0)
+    class2 = (tan_value >= 1) & (u >= 0) & (v >= 0)
+    class3 = (tan_value < -1) & (u <= 0) & (v >= 0)
+    class4 = (tan_value < 0) & (tan_value >= -1) & (u < 0) & (v >= 0)
+    class8 = (tan_value >= -1) & (tan_value < 0) & (u > 0) & (v <= 0)
+    class7 = (tan_value < -1) & (u >= 0) & (v <= 0)
+    class6 = (tan_value >= 1) & (u <= 0) & (v <= 0)
+    class5 = (tan_value >= 0) & (tan_value < 1) & (u < 0) & (v <= 0)
+
+    seg = np.zeros((h, w))
+
+    seg[class1] = 1
+    seg[class2] = 2
+    seg[class3] = 3
+    seg[class4] = 4
+    seg[class5] = 5
+    seg[class6] = 6
+    seg[class7] = 7
+    seg[class8] = 8
+    seg[class0] = 0
+    seg[idx] = 0
+
+    return seg
+
+
+def flow_error(tu, tv, u, v):
+    """
+    Calculate average end point error
+    :param tu: ground-truth horizontal flow map
+    :param tv: ground-truth vertical flow map
+    :param u:  estimated horizontal flow map
+    :param v:  estimated vertical flow map
+    :return: End point error of the estimated flow
+    """
+    smallflow = 0.0
+    """
+    stu = tu[bord+1:end-bord,bord+1:end-bord]
+    stv = tv[bord+1:end-bord,bord+1:end-bord]
+    su = u[bord+1:end-bord,bord+1:end-bord]
+    sv = v[bord+1:end-bord,bord+1:end-bord]
+    """
+    stu = tu[:]
+    stv = tv[:]
+    su = u[:]
+    sv = v[:]
+
+    idxUnknow = (abs(stu) > UNKNOWN_FLOW_THRESH) | (abs(stv) > UNKNOWN_FLOW_THRESH)
+    stu[idxUnknow] = 0
+    stv[idxUnknow] = 0
+    su[idxUnknow] = 0
+    sv[idxUnknow] = 0
+
+    ind2 = [(np.absolute(stu) > smallflow) | (np.absolute(stv) > smallflow)]
+    index_su = su[ind2]
+    index_sv = sv[ind2]
+    an = 1.0 / np.sqrt(index_su ** 2 + index_sv ** 2 + 1)
+    un = index_su * an
+    vn = index_sv * an
+
+    index_stu = stu[ind2]
+    index_stv = stv[ind2]
+    tn = 1.0 / np.sqrt(index_stu ** 2 + index_stv ** 2 + 1)
+    tun = index_stu * tn
+    tvn = index_stv * tn
+
+    """
+    angle = un * tun + vn * tvn + (an * tn)
+    index = [angle == 1.0]
+    angle[index] = 0.999
+    ang = np.arccos(angle)
+    mang = np.mean(ang)
+    mang = mang * 180 / np.pi
+    """
+
+    epe = np.sqrt((stu - su) ** 2 + (stv - sv) ** 2)
+    epe = epe[ind2]
+    return np.mean(epe)
+
+
+def flow_to_image(flow, display=False):
+    """
+    Convert flow into middlebury color code image
+    :param flow: optical flow map
+    :return: optical flow image in middlebury color
+    """
+    u = flow[:, :, 0]
+    v = flow[:, :, 1]
+
+    maxu = -999.0
+    maxv = -999.0
+    minu = 999.0
+    minv = 999.0
+
+    idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
+    u[idxUnknow] = 0
+    v[idxUnknow] = 0
+
+    maxu = max(maxu, np.max(u))
+    minu = min(minu, np.min(u))
+
+    maxv = max(maxv, np.max(v))
+    minv = min(minv, np.min(v))
+
+    rad = np.sqrt(u ** 2 + v ** 2)
+    maxrad = max(-1, np.max(rad))
+
+    if display:
+        print("max flow: %.4f\nflow range:\nu = %.3f .. %.3f\nv = %.3f .. %.3f" % (maxrad, minu, maxu, minv, maxv))
+
+    u = u / (maxrad + np.finfo(float).eps)
+    v = v / (maxrad + np.finfo(float).eps)
+
+    img = compute_color(u, v)
+
+    idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
+    img[idx] = 0
+
+    return np.uint8(img)
+
+
+def evaluate_flow_file(gt, pred):
+    """
+    evaluate the estimated optical flow end point error according to ground truth provided
+    :param gt: ground truth file path
+    :param pred: estimated optical flow file path
+    :return: end point error, float32
+    """
+    # Read flow files and calculate the errors
+    gt_flow = read_flow(gt)  # ground truth flow
+    eva_flow = read_flow(pred)  # predicted flow
+    # Calculate errors
+    average_pe = flow_error(gt_flow[:, :, 0], gt_flow[:, :, 1], eva_flow[:, :, 0], eva_flow[:, :, 1])
+    return average_pe
+
+
+def evaluate_flow(gt_flow, pred_flow):
+    """
+    gt: ground-truth flow
+    pred: estimated flow
+    """
+    return flow_error(
+        gt_flow[:, :, 0],
+        gt_flow[:, :, 1],
+        pred_flow[:, :, 0],
+        pred_flow[:, :, 1],
+    )
+
+
+"""
+==============
+Disparity Section
+==============
+"""
+
+
+def read_disp_png(file_name):
+    """
+    Read optical flow from KITTI .png file
+    :param file_name: name of the flow file
+    :return: optical flow data in matrix
+    """
+    image_object = png.Reader(filename=file_name)
+    image_direct = image_object.asDirect()
+    image_data = list(image_direct[2])
+    (w, h) = image_direct[3]["size"]
+    channel = len(image_data[0]) / w
+    flow = np.zeros((h, w, channel), dtype=np.uint16)
+    for i in range(len(image_data)):
+        for j in range(channel):
+            flow[i, :, j] = image_data[i][j::channel]
+    return flow[:, :, 0] / 256
+
+
+def disp_to_flowfile(disp, filename):
+    """
+    Read KITTI disparity file in png format
+    :param disp: disparity matrix
+    :param filename: the flow file name to save
+    :return: None
+    """
+    f = open(filename, "wb")
+    magic = np.array([202021.25], dtype=np.float32)
+    (height, width) = disp.shape[0:2]
+    w = np.array([width], dtype=np.int32)
+    h = np.array([height], dtype=np.int32)
+    empty_map = np.zeros((height, width), dtype=np.float32)
+    data = np.dstack((disp, empty_map))
+    magic.tofile(f)
+    w.tofile(f)
+    h.tofile(f)
+    data.tofile(f)
+    f.close()
+
+
+"""
+==============
+Image Section
+==============
+"""
+
+
+def read_image(filename):
+    """
+    Read normal image of any format
+    :param filename: name of the image file
+    :return: image data in matrix uint8 type
+    """
+    img = Image.open(filename)
+    return np.array(img)
+
+
+def warp_image(im, flow):
+    """
+    Use optical flow to warp image to the next
+    :param im: image to warp
+    :param flow: optical flow
+    :return: warped image
+    """
+    from scipy import interpolate
+
+    image_height = im.shape[0]
+    image_width = im.shape[1]
+    flow_height = flow.shape[0]
+    flow_width = flow.shape[1]
+    n = image_height * image_width
+    (iy, ix) = np.float32(np.mgrid[0:image_height, 0:image_width])
+    (fy, fx) = np.float32(np.mgrid[0:flow_height, 0:flow_width])
+    fx += flow[:, :, 0].astype(np.float)
+    fy += flow[:, :, 1]
+    mask = np.logical_or(fx < 0, fx > flow_width)
+    mask = np.logical_or(mask, fy < 0)
+    mask = np.logical_or(mask, fy > flow_height)
+    fx = np.minimum(np.maximum(fx, 0), flow_width)
+    fy = np.minimum(np.maximum(fy, 0), flow_height)
+    points = np.concatenate((ix.reshape(n, 1), iy.reshape(n, 1)), axis=1)
+    xi = np.concatenate((fx.reshape(n, 1), fy.reshape(n, 1)), axis=1)
+    warp = np.zeros((image_height, image_width, im.shape[2]))
+    for i in range(im.shape[2]):
+        channel = im[:, :, i]
+        # plt.imshow(channel, cmap='gray')
+        values = channel.reshape(n, 1)
+        new_channel = interpolate.griddata(points, values, xi, method="cubic")
+        new_channel = np.reshape(new_channel, [flow_height, flow_width])
+        new_channel[mask] = 1
+        warp[:, :, i] = new_channel.astype(np.uint8)
+
+    return warp.astype(np.uint8)
+
+
+"""
+==============
+Others
+==============
+"""
+
+
+def scale_image(image, new_range):
+    """
+    Linearly scale the image into desired range
+    :param image: input image
+    :param new_range: the new range to be aligned
+    :return: image normalized in new range
+    """
+    min_val = np.min(image).astype(np.float32)
+    max_val = np.max(image).astype(np.float32)
+    min_val_new = np.array(min(new_range), dtype=np.float32)
+    max_val_new = np.array(max(new_range), dtype=np.float32)
+    scaled_image = (image - min_val) / (max_val - min_val) * (max_val_new - min_val_new) + min_val_new
+    return scaled_image.astype(np.uint8)
+
+
+def compute_color(u, v):
+    """
+    compute optical flow color map
+    :param u: optical flow horizontal map
+    :param v: optical flow vertical map
+    :return: optical flow in color code
+    """
+    [h, w] = u.shape
+    img = np.zeros([h, w, 3])
+    nanIdx = np.isnan(u) | np.isnan(v)
+    u[nanIdx] = 0
+    v[nanIdx] = 0
+
+    colorwheel = make_color_wheel()
+    ncols = np.size(colorwheel, 0)
+
+    rad = np.sqrt(u ** 2 + v ** 2)
+
+    a = np.arctan2(-v, -u) / np.pi
+
+    fk = (a + 1) / 2 * (ncols - 1) + 1
+
+    k0 = np.floor(fk).astype(int)
+
+    k1 = k0 + 1
+    k1[k1 == ncols + 1] = 1
+    f = fk - k0
+
+    for i in range(np.size(colorwheel, 1)):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0 - 1] / 255
+        col1 = tmp[k1 - 1] / 255
+        col = (1 - f) * col0 + f * col1
+
+        idx = rad <= 1
+        col[idx] = 1 - rad[idx] * (1 - col[idx])
+        notidx = np.logical_not(idx)
+
+        col[notidx] *= 0.75
+        img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx)))
+
+    return img
+
+
+def make_color_wheel():
+    """
+    Generate color wheel according Middlebury color code
+    :return: Color wheel
+    """
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+
+    colorwheel = np.zeros([ncols, 3])
+
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))
+    col += RY
+
+    # YG
+    colorwheel[col : col + YG, 0] = 255 - np.transpose(np.floor(255 * np.arange(0, YG) / YG))
+    colorwheel[col : col + YG, 1] = 255
+    col += YG
+
+    # GC
+    colorwheel[col : col + GC, 1] = 255
+    colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC))
+    col += GC
+
+    # CB
+    colorwheel[col : col + CB, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, CB) / CB))
+    colorwheel[col : col + CB, 2] = 255
+    col += CB
+
+    # BM
+    colorwheel[col : col + BM, 2] = 255
+    colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM))
+    col += +BM
+
+    # MR
+    colorwheel[col : col + MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+    colorwheel[col : col + MR, 0] = 255
+
+    return colorwheel

utils/tb_image_recorder.py

@@ -0,0 +1,28 @@
+import threading
+
+import numpy as np
+
+
+class TBImageRecorder(threading.Thread):
+    """
+    TBImageRecorder
+    """
+
+    def __init__(self, tb_writer, func, queue):
+        super(TBImageRecorder, self).__init__()
+        self._tb_writer = tb_writer
+        self._func = func
+        self._queue = queue
+
+    def run(self):
+        while True:
+            msgs, iter_index = self._queue.get()
+            if msgs:
+                img_info = self._func(*msgs)
+
+                print("logging the images")
+                for tag, images in img_info.items():
+                    if images is not None:
+                        self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index)
+            else:
+                break

utils/util_tensorboard.py

@@ -0,0 +1,43 @@
+import threading
+
+import numpy as np
+
+from utils.util import to_np
+
+
+def histogram_logger(tb_writer, iter_index, net=None):
+    if net is not None:
+        for tag, value in net.named_parameters():
+            tag = tag.replace('.', '/')
+            tb_writer.add_histogram(tag, to_np(value), iter_index)
+            tb_writer.add_histogram(tag + '/grad', to_np(value.grad), iter_index)
+
+
+def value_logger(tb_writer, iter_index, loss_info):
+    for tag, value in loss_info.items():
+        tb_writer.add_scalar(tag, value, iter_index)
+
+
+class TBImageRecorder(threading.Thread):
+    """
+    TBImageRecorder
+    """
+
+    def __init__(self, tb_writer, func, queue):
+        super(TBImageRecorder, self).__init__()
+        self._tb_writer = tb_writer
+        self._func = func
+        self._queue = queue
+
+    def run(self):
+        while True:
+            msgs, iter_index = self._queue.get()
+            if msgs:
+                img_info = self._func(*msgs)
+
+                print("logging the images")
+                for tag, images in img_info.items():
+                    if images is not None:
+                        self._tb_writer.add_image(tag, np.clip(images, 0, 255).astype(np.uint8), iter_index)
+            else:
+                break

utils/warping.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_grid(x):
+    torchHorizontal = (
+        torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3))
+    )
+    torchVertical = (
+        torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3))
+    )
+    return torch.cat([torchHorizontal, torchVertical], 1).cuda()
+
+
+class WarpingLayer(nn.Module):
+    def __init__(self, device):
+        super(WarpingLayer, self).__init__()
+        self.device = device
+
+    def forward(self, x, flow):
+        # WarpingLayer uses F.grid_sample, which expects normalized grid
+        # we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code
+        # so here we need to denormalize the flow
+        flow_for_grip = torch.zeros_like(flow).cuda()
+        flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0)
+        flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0)
+
+        grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1)
+        return F.grid_sample(x, grid, align_corners=True)

warping.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_grid(x):
+    torchHorizontal = (
+        torch.linspace(-1.0, 1.0, x.size(3)).view(1, 1, 1, x.size(3)).expand(x.size(0), 1, x.size(2), x.size(3))
+    )
+    torchVertical = (
+        torch.linspace(-1.0, 1.0, x.size(2)).view(1, 1, x.size(2), 1).expand(x.size(0), 1, x.size(2), x.size(3))
+    )
+    return torch.cat([torchHorizontal, torchVertical], 1).cuda()
+
+
+class WarpingLayer(nn.Module):
+    def __init__(self, device):
+        super(WarpingLayer, self).__init__()
+        self.device = device
+
+    def forward(self, x, flow):
+        # WarpingLayer uses F.VGG19_pytorch, which expects normalized grid
+        # we still output unnormalized flow for the convenience of comparing EPEs with FlowNet2 and original code
+        # so here we need to denormalize the flow
+        flow_for_grip = torch.zeros_like(flow).cuda()
+        flow_for_grip[:, 0, :, :] = flow[:, 0, :, :] / ((flow.size(3) - 1.0) / 2.0)
+        flow_for_grip[:, 1, :, :] = flow[:, 1, :, :] / ((flow.size(2) - 1.0) / 2.0)
+
+        grid = (get_grid(x) + flow_for_grip).permute(0, 2, 3, 1)
+        return F.VGG19_pytorch(x, grid, align_corners=True)