Add application file

app.py

@@ -0,0 +1,81 @@
+import os.path
+
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+import test
+from PIL import Image
+import gdown
+
+
+st.set_page_config(layout="wide")
+
+# Specify canvas parameters in application
+drawing_object = st.sidebar.selectbox(
+    "Object:", ("sea", "cloud", "bush", "grass", "mountain", "sky", "snow",
+                "tree", "flower", "road")
+)
+drawing_object_dict = {"sea": "rgb(56,79,131)", "cloud": "rgb(239,239,239)",
+                       "bush": "rgb(93,110,50)", "grass": "rgb(183,210,78)",
+                       "mountain": "rgb(60,59,75)", "snow": "rgb(250,250,250)",
+                       "sky": "rgb(117,158,223)", "tree": "rgb(53, 38, 19)",
+                       "flower": "rgb(230,112,182)",
+                       "road": "rgb(152, 126, 106)"}
+
+stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 3)
+
+stroke_color = drawing_object_dict[drawing_object]
+
+
+col1, col2 = st.columns(2)
+with col1:
+    # Create a canvas component with different parameters
+    canvas_result = st_canvas(
+        fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
+        stroke_width=stroke_width,
+        stroke_color=stroke_color,
+        background_color="rgb(117,158,223)",
+        background_image=None,
+        height=512,
+        width=512,
+        drawing_mode="freedraw",
+        point_display_radius=0,
+        key="canvas",
+    )
+    if canvas_result.image_data is not None:
+        pass
+
+
+@st.cache
+def download_model():
+    f_checkpoint = os.path.join("latest_net_G.pth")
+    if not os.path.exists(f_checkpoint):
+        with st.spinner("Downloading model... this may take awhile! \n Don't stop it!"):
+            url = 'https://drive.google.com/uc?id=15VSa2m2F6Ch0NpewDR7mkKAcXlMgDi5F'
+            output = 'latest_net_G.pth'
+            gdown.download(url, output, quiet=False)
+
+
+if st.button('generate'):
+    download_model()
+    image = Image.fromarray(canvas_result.image_data)
+    s = test.semantic(image)
+    image = test.evaluate(s)
+    image = test.to_image(image)
+    with col2:
+        st.image(image, clamp=True, width=512)
+
+
+st.markdown(
+    """
+    <style>
+    [data-testid="stSidebar"][aria-expanded="true"] > div:first-child {
+        width: 120px;
+    }
+    [data-testid="stSidebar"][aria-expanded="false"] > div:first-child {
+        width: 500px;
+        margin-left: -500px;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True,
+)

label_colors.py

@@ -0,0 +1,12 @@
+colorMap = [
+    {"color": (56, 79, 131), "id": 154, "label": "sea"},
+    {"color": (239, 239, 239), "id": 105, "label": "cloud"},
+    {"color": (93, 110, 50), "id": 96, "label": "bush"},
+    {"color": (183, 210, 78), "id": 123, "label": "grass"},
+    {"color": (60, 59, 75), "id": 134, "label": "mountain"},
+    {"color": (117, 158, 223), "id": 156, "label": "sky"},
+    {"color": (250, 250, 250), "id": 158, "label": "snow"},
+    {"color": (53, 38, 19), "id": 168, "label": "tree"},
+    {"color": (230, 112, 182), "id": 118, "label": "flower"},
+    {"color": (152, 126, 106), "id": 148, "label": "road"}
+]

spade/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Jiayuan MAO
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

spade/README.md

@@ -0,0 +1,118 @@
+# Synchronized-BatchNorm-PyTorch
+
+**IMPORTANT: Please read the "Implementation details and highlights" section before use.**
+
+Synchronized Batch Normalization implementation in PyTorch.
+
+This module differs from the built-in PyTorch BatchNorm as the mean and
+standard-deviation are reduced across all devices during training.
+
+For example, when one uses `nn.DataParallel` to wrap the network during
+training, PyTorch's implementation normalize the tensor on each device using
+the statistics only on that device, which accelerated the computation and
+is also easy to implement, but the statistics might be inaccurate.
+Instead, in this synchronized version, the statistics will be computed
+over all training samples distributed on multiple devices.
+
+Note that, for one-GPU or CPU-only case, this module behaves exactly same
+as the built-in PyTorch implementation.
+
+This module is currently only a prototype version for research usages. As mentioned below,
+it has its limitations and may even suffer from some design problems. If you have any
+questions or suggestions, please feel free to
+[open an issue](https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues) or 
+[submit a pull request](https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues).
+
+## Why Synchronized BatchNorm?
+
+Although the typical implementation of BatchNorm working on multiple devices (GPUs)
+is fast (with no communication overhead), it inevitably reduces the size of batch size,
+which potentially degenerates the performance. This is not a significant issue in some
+standard vision tasks such as ImageNet classification (as the batch size per device
+is usually large enough to obtain good statistics). However, it will hurt the performance
+in some tasks that the batch size is usually very small (e.g., 1 per GPU).
+
+For example, the importance of synchronized batch normalization in object detection has been recently proved with a
+an extensive analysis in the paper [MegDet: A Large Mini-Batch Object Detector](https://arxiv.org/abs/1711.07240).
+
+## Usage
+
+To use the Synchronized Batch Normalization, we add a data parallel replication callback. This introduces a slight
+difference with typical usage of the `nn.DataParallel`.
+
+Use it with a provided, customized data parallel wrapper:
+
+```python
+from sync_batchnorm import SynchronizedBatchNorm1d, DataParallelWithCallback
+
+sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+```
+
+Or, if you are using a customized data parallel module, you can use this library as a monkey patching.
+
+```python
+from torch.nn import DataParallel  # or your customized DataParallel module
+from sync_batchnorm import SynchronizedBatchNorm1d, patch_replication_callback
+
+sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+patch_replication_callback(sync_bn)  # monkey-patching
+```
+
+You can use `convert_model` to convert your model to use Synchronized BatchNorm easily.
+
+```python
+import torch.nn as nn
+from torchvision import models
+from sync_batchnorm import convert_model
+# m is a standard pytorch model
+m = models.resnet18(True)
+m = nn.DataParallel(m)
+# after convert, m is using SyncBN
+m = convert_model(m)
+```
+
+See also `tests/test_sync_batchnorm.py` for numeric result comparison.
+
+## Implementation details and highlights
+
+If you are interested in how batch statistics are reduced and broadcasted among multiple devices, please take a look
+at the code with detailed comments. Here we only emphasize some highlights of the implementation:
+
+- This implementation is in pure-python. No C++ extra extension libs.
+- Easy to use as demonstrated above.
+- It uses unbiased variance to update the moving average, and use `sqrt(max(var, eps))` instead of `sqrt(var + eps)`.
+- The implementation requires that each module on different devices should invoke the `batchnorm` for exactly SAME
+amount of times in each forward pass. For example, you can not only call `batchnorm` on GPU0 but not on GPU1. The `#i
+(i = 1, 2, 3, ...)` calls of the `batchnorm` on each device will be viewed as a whole and the statistics will be reduced.
+This is tricky but is a good way to handle PyTorch's dynamic computation graph. Although sounds complicated, this
+will usually not be the issue for most of the models.
+
+## Known issues
+
+#### Runtime error on backward pass.
+
+Due to a [PyTorch Bug](https://github.com/pytorch/pytorch/issues/3883), using old PyTorch libraries will trigger an `RuntimeError` with messages like:
+
+```
+Assertion `pos >= 0 && pos < buffer.size()` failed.
+```
+
+This has already been solved in the newest PyTorch repo, which, unfortunately, has not been pushed to the official and anaconda binary release. Thus, you are required to build the PyTorch package from the source according to the
+ instructions [here](https://github.com/pytorch/pytorch#from-source).
+
+#### Numeric error.
+
+Because this library does not fuse the normalization and statistics operations in C++ (nor CUDA), it is less
+numerically stable compared to the original PyTorch implementation. Detailed analysis can be found in
+`tests/test_sync_batchnorm.py`.
+
+## Authors and License:
+
+Copyright (c) 2018-, [Jiayuan Mao](https://vccy.xyz).
+
+**Contributors**: [Tete Xiao](https://tetexiao.com), [DTennant](https://github.com/DTennant).
+
+Distributed under **MIT License** (See LICENSE)
+

spade/Synchronized-BatchNorm-PyTorch

@@ -0,0 +1 @@
+Subproject commit dcfae91cbc3767a3c5cd28d46ab78503a22b0fe7

spade/dataset.py

@@ -0,0 +1,24 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torchvision.transforms as transforms
+from PIL import Image
+
+def __scale_width(img, target_width, method=Image.BICUBIC):
+  ow, oh = img.size
+  if (ow == target_width):
+    return img
+  w = target_width
+  h = int(target_width * oh / ow)
+  return img.resize((w, h), method)
+
+def get_transform(opt, method=Image.BICUBIC, normalize=True):
+  transform_list = []
+  transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt['load_size'], method)))
+  transform_list += [transforms.ToTensor()]
+  if normalize:
+    transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+  
+  return transforms.Compose(transform_list)

spade/generator.py

@@ -0,0 +1,113 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from spade.normalizer import SPADE
+
+class SPADEGenerator(nn.Module):
+  def __init__(self, opt):
+    super().__init__()
+
+    # nf: # of gen filters in first conv layer
+    nf = 64
+
+    self.sw, self.sh = self.compute_latent_vector_size(opt['crop_size'], opt['aspect_ratio'])
+
+    self.fc = nn.Conv2d(opt['label_nc'], 16 * nf, 3, padding=1)
+
+    self.head_0 = SPADEResnetBlock(opt, 16 * nf, 16 * nf)
+
+    self.G_middle_0 = SPADEResnetBlock(opt, 16 * nf, 16 * nf)
+    self.G_middle_1 = SPADEResnetBlock(opt, 16 * nf, 16 * nf)
+
+    self.up_0 = SPADEResnetBlock(opt, 16 * nf, 8 * nf)
+    self.up_1 = SPADEResnetBlock(opt, 8 * nf, 4 * nf)
+    self.up_2 = SPADEResnetBlock(opt, 4 * nf, 2 * nf)
+    self.up_3 = SPADEResnetBlock(opt, 2 * nf, 1 * nf)
+
+    self.conv_img = nn.Conv2d(1 * nf, 3, 3, padding=1)
+
+    self.up = nn.Upsample(scale_factor=2)
+  
+  def compute_latent_vector_size(self, crop_size, aspect_ratio):
+    num_up_layers = 5
+
+    sw = crop_size // (2**num_up_layers)
+    sh = round(sw / aspect_ratio)
+
+    return sw, sh
+  
+  def forward(self, seg):
+    # we downsample segmap and run convolution
+    x = F.interpolate(seg, size=(self.sh, self.sw))
+    x = self.fc(x)
+
+    x = self.head_0(x, seg)
+
+    x = self.up(x)
+    x = self.G_middle_0(x, seg)
+    x = self.G_middle_1(x, seg)
+
+    x = self.up(x)
+    x = self.up_0(x, seg)
+    x = self.up(x)
+    x = self.up_1(x, seg)
+    x = self.up(x)
+    x = self.up_2(x, seg)
+    x = self.up(x)
+    x = self.up_3(x, seg)
+
+    x = self.conv_img(F.leaky_relu(x, 2e-1))
+    x = torch.tanh(x)
+
+    return x
+
+import torch.nn.utils.spectral_norm as spectral_norm
+
+# label_nc: the #channels of the input semantic map, hence the input dim of SPADE
+# label_nc: also equivalent to the # of input label classes
+class SPADEResnetBlock(nn.Module):
+  def __init__(self, opt, fin, fout):
+    super().__init__()
+
+    self.learned_shortcut = (fin != fout)
+    fmiddle = min(fin, fout)
+
+    self.conv_0 = spectral_norm(nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1))
+    self.conv_1 = spectral_norm(nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1))
+    if self.learned_shortcut:
+      self.conv_s = spectral_norm(nn.Conv2d(fin, fout, kernel_size=1, bias=False))
+
+    # define normalization layers
+    self.norm_0 = SPADE(opt, fin)
+    self.norm_1 = SPADE(opt, fmiddle)
+    if self.learned_shortcut:
+      self.norm_s = SPADE(opt, fin)
+
+  # note the resnet block with SPADE also takes in |seg|,
+  # the semantic segmentation map as input
+  def forward(self, x, seg):
+    x_s = self.shortcut(x, seg)
+
+    dx = self.conv_0(self.relu(self.norm_0(x, seg)))
+    dx = self.conv_1(self.relu(self.norm_1(dx, seg)))
+
+    out = x_s + dx
+    return out
+
+  def shortcut(self, x, seg):
+    if self.learned_shortcut:
+      x_s = self.conv_s(self.norm_s(x, seg))
+    else:
+      x_s = x
+    return x_s
+
+  def relu(self, x):
+    return F.leaky_relu(x, 2e-1)
+
+

spade/model.py

@@ -0,0 +1,101 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import os
+import torch
+from torch.nn import init
+
+from spade.generator import SPADEGenerator
+
+
+class Pix2PixModel(torch.nn.Module):
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        self.FloatTensor = torch.cuda.FloatTensor if opt['use_gpu'] \
+            else torch.FloatTensor
+
+        self.netG = self.initialize_networks(opt)
+
+    def forward(self, data, mode):
+        input_semantics, real_image = self.preprocess_input(data)
+
+        if mode == 'inference':
+            with torch.no_grad():
+                fake_image = self.generate_fake(input_semantics)
+            return fake_image
+        else:
+            raise ValueError("|mode| is invalid")
+
+    def preprocess_input(self, data):
+        data['label'] = data['label'].long()
+
+        # move to GPU and change data types
+        if self.opt['use_gpu']:
+            data['label'] = data['label'].cuda()
+            data['instance'] = data['instance'].cuda()
+            data['image'] = data['image'].cuda()
+
+        # create one-hot label map
+        label_map = data['label']
+        bs, _, h, w = label_map.size()
+        input_label = self.FloatTensor(bs, self.opt['label_nc'], h, w).zero_()
+        # one whole label map -> to one label map per class
+        input_semantics = input_label.scatter_(1, label_map, 1.0)
+
+        return input_semantics, data['image']
+
+    def generate_fake(self, input_semantics):
+        fake_image = self.netG(input_semantics)
+        return fake_image
+
+    def create_network(self, cls, opt):
+        net = cls(opt)
+        if self.opt['use_gpu']:
+            net.cuda()
+
+        gain = 0.02
+
+        def init_weights(m):
+            classname = m.__class__.__name__
+            if classname.find('BatchNorm2d') != -1:
+                if hasattr(m, 'weight') and m.weight is not None:
+                    init.normal_(m.weight.data, 1.0, gain)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+            elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+                init.xavier_normal_(m.weight.data, gain=gain)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+
+        # Applies fn recursively to every submodule (as returned by .children()) as well as self
+        net.apply(init_weights)
+
+        return net
+
+    def load_network(self, net, label, epoch, opt):
+        save_filename = '%s_net_%s.pth' % (epoch, label)
+        save_path = os.path.join( save_filename)
+        weights = torch.load(save_path)
+        net.load_state_dict(weights)
+        return net
+
+    def initialize_networks(self, opt):
+        netG = self.create_network(SPADEGenerator, opt)
+
+        if not opt['isTrain']:
+            netG = self.load_network(netG, 'G', opt['which_epoch'], opt)
+
+        # self.print_network(netG)
+
+        return netG
+
+    def print_network(self, net):
+        num_params = 0
+        for param in net.parameters():
+            num_params += param.numel()
+        print('Network [%s] was created. Total number of parameters: %.1f million. '
+              % (type(net).__name__, num_params / 1000000))
+        print(net)

spade/normalizer.py

@@ -0,0 +1,49 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+# norm_nc: the #channels of the normalized activations, hence the output dim of SPADE
+# label_nc: the #channels of the input semantic map, hence the input dim of SPADE
+# label_nc: also equivalent to the # of input label classes
+class SPADE(nn.Module):
+  def __init__(self, opt, norm_nc):
+    super().__init__()
+
+    self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False)
+
+    # number of internal filters for generating scale/bias
+    nhidden = 128
+    # size of kernels
+    kernal_size = 3
+    # padding size
+    padding = kernal_size // 2
+
+    self.mlp_shared = nn.Sequential(
+      nn.Conv2d(opt['label_nc'], nhidden, kernel_size=kernal_size, padding=padding),
+      nn.ReLU()
+    )
+    self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=kernal_size, padding=padding)
+    self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=kernal_size, padding=padding)
+
+  def forward(self, x, segmap):
+    # Part 1. generate parameter-free normalized activations
+    normalized = self.param_free_norm(x)
+
+    # Part 2. produce scaling and bias conditioned on semantic map
+    # resize input segmentation map to match x.size() using nearest interpolation
+    # N, C, H, W = x.size()
+    segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+    actv = self.mlp_shared(segmap)
+    gamma = self.mlp_gamma(actv)
+    beta = self.mlp_beta(actv)
+
+    # apply scale and bias
+    out = normalized * (1 + gamma) + beta
+
+    return out

spade/tests/test_numeric_batchnorm.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# File   : test_numeric_batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+
+import unittest
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+from sync_batchnorm.unittest import TorchTestCase
+
+
+def handy_var(a, unbias=True):
+    n = a.size(0)
+    asum = a.sum(dim=0)
+    as_sum = (a ** 2).sum(dim=0)  # a square sum
+    sumvar = as_sum - asum * asum / n
+    if unbias:
+        return sumvar / (n - 1)
+    else:
+        return sumvar / n
+
+
+class NumericTestCase(TorchTestCase):
+    def testNumericBatchNorm(self):
+        a = torch.rand(16, 10)
+        bn = nn.BatchNorm1d(10, momentum=1, eps=1e-5, affine=False)
+        bn.train()
+
+        a_var1 = Variable(a, requires_grad=True)
+        b_var1 = bn(a_var1)
+        loss1 = b_var1.sum()
+        loss1.backward()
+
+        a_var2 = Variable(a, requires_grad=True)
+        a_mean2 = a_var2.mean(dim=0, keepdim=True)
+        a_std2 = torch.sqrt(handy_var(a_var2, unbias=False).clamp(min=1e-5))
+        # a_std2 = torch.sqrt(a_var2.var(dim=0, keepdim=True, unbiased=False) + 1e-5)
+        b_var2 = (a_var2 - a_mean2) / a_std2
+        loss2 = b_var2.sum()
+        loss2.backward()
+
+        self.assertTensorClose(bn.running_mean, a.mean(dim=0))
+        self.assertTensorClose(bn.running_var, handy_var(a))
+        self.assertTensorClose(a_var1.data, a_var2.data)
+        self.assertTensorClose(b_var1.data, b_var2.data)
+        self.assertTensorClose(a_var1.grad, a_var2.grad)
+
+
+if __name__ == '__main__':
+    unittest.main()

spade/tests/test_numeric_batchnorm_v2.py

@@ -0,0 +1,62 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : test_numeric_batchnorm_v2.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 11/01/2018
+#
+# Distributed under terms of the MIT license.
+
+"""
+Test the numerical implementation of batch normalization.
+
+Author: acgtyrant.
+See also: https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+"""
+
+import unittest
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from sync_batchnorm.unittest import TorchTestCase
+from sync_batchnorm.batchnorm_reimpl import BatchNorm2dReimpl
+
+
+class NumericTestCasev2(TorchTestCase):
+    def testNumericBatchNorm(self):
+        CHANNELS = 16
+        batchnorm1 = nn.BatchNorm2d(CHANNELS, momentum=1)
+        optimizer1 = optim.SGD(batchnorm1.parameters(), lr=0.01)
+
+        batchnorm2 = BatchNorm2dReimpl(CHANNELS, momentum=1)
+        batchnorm2.weight.data.copy_(batchnorm1.weight.data)
+        batchnorm2.bias.data.copy_(batchnorm1.bias.data)
+        optimizer2 = optim.SGD(batchnorm2.parameters(), lr=0.01)
+
+        for _ in range(100):
+            input_ = torch.rand(16, CHANNELS, 16, 16)
+
+            input1 = input_.clone().requires_grad_(True)
+            output1 = batchnorm1(input1)
+            output1.sum().backward()
+            optimizer1.step()
+
+            input2 = input_.clone().requires_grad_(True)
+            output2 = batchnorm2(input2)
+            output2.sum().backward()
+            optimizer2.step()
+
+        self.assertTensorClose(input1, input2)
+        self.assertTensorClose(output1, output2)
+        self.assertTensorClose(input1.grad, input2.grad)
+        self.assertTensorClose(batchnorm1.weight.grad, batchnorm2.weight.grad)
+        self.assertTensorClose(batchnorm1.bias.grad, batchnorm2.bias.grad)
+        self.assertTensorClose(batchnorm1.running_mean, batchnorm2.running_mean)
+        self.assertTensorClose(batchnorm2.running_mean, batchnorm2.running_mean)
+
+
+if __name__ == '__main__':
+    unittest.main()
+

spade/tests/test_sync_batchnorm.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+# File   : test_sync_batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+
+import unittest
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+from sync_batchnorm import set_sbn_eps_mode
+from sync_batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, DataParallelWithCallback
+from sync_batchnorm.unittest import TorchTestCase
+
+set_sbn_eps_mode('plus')
+
+
+def handy_var(a, unbias=True):
+    n = a.size(0)
+    asum = a.sum(dim=0)
+    as_sum = (a ** 2).sum(dim=0)  # a square sum
+    sumvar = as_sum - asum * asum / n
+    if unbias:
+        return sumvar / (n - 1)
+    else:
+        return sumvar / n
+
+
+def _find_bn(module):
+    for m in module.modules():
+        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, SynchronizedBatchNorm1d, SynchronizedBatchNorm2d)):
+            return m
+
+
+class SyncTestCase(TorchTestCase):
+    def _syncParameters(self, bn1, bn2):
+        bn1.reset_parameters()
+        bn2.reset_parameters()
+        if bn1.affine and bn2.affine:
+            bn2.weight.data.copy_(bn1.weight.data)
+            bn2.bias.data.copy_(bn1.bias.data)
+
+    def _checkBatchNormResult(self, bn1, bn2, input, is_train, cuda=False):
+        """Check the forward and backward for the customized batch normalization."""
+        bn1.train(mode=is_train)
+        bn2.train(mode=is_train)
+
+        if cuda:
+            input = input.cuda()
+
+        self._syncParameters(_find_bn(bn1), _find_bn(bn2))
+
+        input1 = Variable(input, requires_grad=True)
+        output1 = bn1(input1)
+        output1.sum().backward()
+        input2 = Variable(input, requires_grad=True)
+        output2 = bn2(input2)
+        output2.sum().backward()
+
+        self.assertTensorClose(input1.data, input2.data)
+        self.assertTensorClose(output1.data, output2.data)
+        self.assertTensorClose(input1.grad, input2.grad)
+        self.assertTensorClose(_find_bn(bn1).running_mean, _find_bn(bn2).running_mean)
+        self.assertTensorClose(_find_bn(bn1).running_var, _find_bn(bn2).running_var)
+
+    def testSyncBatchNormNormalTrain(self):
+        bn = nn.BatchNorm1d(10)
+        sync_bn = SynchronizedBatchNorm1d(10)
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), True)
+
+    def testSyncBatchNormNormalEval(self):
+        bn = nn.BatchNorm1d(10)
+        sync_bn = SynchronizedBatchNorm1d(10)
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), False)
+
+    def testSyncBatchNormSyncTrain(self):
+        bn = nn.BatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+
+        bn.cuda()
+        sync_bn.cuda()
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), True, cuda=True)
+
+    def testSyncBatchNormSyncEval(self):
+        bn = nn.BatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+
+        bn.cuda()
+        sync_bn.cuda()
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), False, cuda=True)
+
+    def testSyncBatchNorm2DSyncTrain(self):
+        bn = nn.BatchNorm2d(10)
+        sync_bn = SynchronizedBatchNorm2d(10)
+        sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+
+        bn.cuda()
+        sync_bn.cuda()
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10, 16, 16), True, cuda=True)
+
+
+if __name__ == '__main__':
+    unittest.main()

sync_batchnorm/__init__.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import set_sbn_eps_mode
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .batchnorm import patch_sync_batchnorm, convert_model
+from .replicate import DataParallelWithCallback, patch_replication_callback

sync_batchnorm/batchnorm.py

@@ -0,0 +1,412 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+import contextlib
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+
+try:
+    from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+except ImportError:
+    ReduceAddCoalesced = Broadcast = None
+
+try:
+    from jactorch.parallel.comm import SyncMaster
+    from jactorch.parallel.data_parallel import JacDataParallel as DataParallelWithCallback
+except ImportError:
+    from .comm import SyncMaster
+    from .replicate import DataParallelWithCallback
+
+__all__ = [
+    'set_sbn_eps_mode',
+    'SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d',
+    'patch_sync_batchnorm', 'convert_model'
+]
+
+
+SBN_EPS_MODE = 'clamp'
+
+
+def set_sbn_eps_mode(mode):
+    global SBN_EPS_MODE
+    assert mode in ('clamp', 'plus')
+    SBN_EPS_MODE = mode
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dimensions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
+        assert ReduceAddCoalesced is not None, 'Can not use Synchronized Batch Normalization without CUDA support.'
+
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine,
+                                                     track_running_stats=track_running_stats)
+
+        if not self.track_running_stats:
+            import warnings
+            warnings.warn('track_running_stats=False is not supported by the SynchronizedBatchNorm.')
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            return F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        assert input.size(1) == self.num_features, 'Channel size mismatch: got {}, expect {}.'.format(input.size(1), self.num_features)
+        input = input.view(input.size(0), self.num_features, -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+
+        # Reduce-and-broadcast the statistics.
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+
+        # Compute the output.
+        if self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        if hasattr(torch, 'no_grad'):
+            with torch.no_grad():
+                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+        else:
+            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+
+        if SBN_EPS_MODE == 'clamp':
+            return mean, bias_var.clamp(self.eps) ** -0.5
+        elif SBN_EPS_MODE == 'plus':
+            return mean, (bias_var + self.eps) ** -0.5
+        else:
+            raise ValueError('Unknown EPS mode: {}.'.format(SBN_EPS_MODE))
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+
+
+@contextlib.contextmanager
+def patch_sync_batchnorm():
+    import torch.nn as nn
+
+    backup = nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d
+
+    nn.BatchNorm1d = SynchronizedBatchNorm1d
+    nn.BatchNorm2d = SynchronizedBatchNorm2d
+    nn.BatchNorm3d = SynchronizedBatchNorm3d
+
+    yield
+
+    nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d = backup
+
+
+def convert_model(module):
+    """Traverse the input module and its child recursively
+       and replace all instance of torch.nn.modules.batchnorm.BatchNorm*N*d
+       to SynchronizedBatchNorm*N*d
+
+    Args:
+        module: the input module needs to be convert to SyncBN model
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> import torchvision
+        >>> # m is a standard pytorch model
+        >>> m = torchvision.models.resnet18(True)
+        >>> m = nn.DataParallel(m)
+        >>> # after convert, m is using SyncBN
+        >>> m = convert_model(m)
+    """
+    if isinstance(module, torch.nn.DataParallel):
+        mod = module.module
+        mod = convert_model(mod)
+        mod = DataParallelWithCallback(mod, device_ids=module.device_ids)
+        return mod
+
+    mod = module
+    for pth_module, sync_module in zip([torch.nn.modules.batchnorm.BatchNorm1d,
+                                        torch.nn.modules.batchnorm.BatchNorm2d,
+                                        torch.nn.modules.batchnorm.BatchNorm3d],
+                                       [SynchronizedBatchNorm1d,
+                                        SynchronizedBatchNorm2d,
+                                        SynchronizedBatchNorm3d]):
+        if isinstance(module, pth_module):
+            mod = sync_module(module.num_features, module.eps, module.momentum, module.affine)
+            mod.running_mean = module.running_mean
+            mod.running_var = module.running_var
+            if module.affine:
+                mod.weight.data = module.weight.data.clone().detach()
+                mod.bias.data = module.bias.data.clone().detach()
+
+    for name, child in module.named_children():
+        mod.add_module(name, convert_model(child))
+
+    return mod

sync_batchnorm/batchnorm_reimpl.py

@@ -0,0 +1,74 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : batchnorm_reimpl.py
+# Author : acgtyrant
+# Date   : 11/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+
+__all__ = ['BatchNorm2dReimpl']
+
+
+class BatchNorm2dReimpl(nn.Module):
+    """
+    A re-implementation of batch normalization, used for testing the numerical
+    stability.
+
+    Author: acgtyrant
+    See also:
+    https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+    """
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super().__init__()
+
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = nn.Parameter(torch.empty(num_features))
+        self.bias = nn.Parameter(torch.empty(num_features))
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_running_stats(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        init.uniform_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, input_):
+        batchsize, channels, height, width = input_.size()
+        numel = batchsize * height * width
+        input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel)
+        sum_ = input_.sum(1)
+        sum_of_square = input_.pow(2).sum(1)
+        mean = sum_ / numel
+        sumvar = sum_of_square - sum_ * mean
+
+        self.running_mean = (
+                (1 - self.momentum) * self.running_mean
+                + self.momentum * mean.detach()
+        )
+        unbias_var = sumvar / (numel - 1)
+        self.running_var = (
+                (1 - self.momentum) * self.running_var
+                + self.momentum * unbias_var.detach()
+        )
+
+        bias_var = sumvar / numel
+        inv_std = 1 / (bias_var + self.eps).pow(0.5)
+        output = (
+                (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) *
+                self.weight.unsqueeze(1) + self.bias.unsqueeze(1))
+
+        return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous()
+

sync_batchnorm/comm.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+
+        Args:
+            identifier: an identifier, usually is the device id.
+
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+
+        Returns: the message to be sent back to the master device.
+
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

sync_batchnorm/replicate.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+import torch
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, x, y):
+        adiff = float((x - y).abs().max())
+        if (y == 0).all():
+            rdiff = 'NaN'
+        else:
+            rdiff = float((adiff / y).abs().max())
+
+        message = (
+            'Tensor close check failed\n'
+            'adiff={}\n'
+            'rdiff={}\n'
+        ).format(adiff, rdiff)
+        self.assertTrue(torch.allclose(x, y, atol=1e-5, rtol=1e-3), message)
+

test.py

@@ -0,0 +1,63 @@
+import numpy as np
+from label_colors import colorMap
+from PIL import Image
+from spade.model import Pix2PixModel
+from spade.dataset import get_transform
+from torchvision.transforms import ToPILImage
+
+'''colors = np.array([[56, 79, 131], [239, 239, 239],
+                   [93, 110, 50], [183, 210, 78],
+                   [60, 59, 75], [250, 250, 250]])'''
+colors = [key['color'] for key in colorMap]
+id_list = [key['id'] for key in colorMap]
+
+
+def semantic(img):
+    print("semantic", type(img))
+    h, w = img.size
+    imrgb = img.convert("RGB")
+    pix = list(imrgb.getdata())
+    mask = [id_list[colors.index(i)] if i in colors else 156 for i in pix]
+    return np.array(mask).reshape(h, w)
+
+
+def evaluate(labelmap):
+    opt = {
+        'label_nc': 182,  # num classes in coco model
+        'crop_size': 512,
+        'load_size': 512,
+        'aspect_ratio': 1.0,
+        'isTrain': False,
+        'checkpoints_dir': 'app',
+        'which_epoch': 'latest',
+        'use_gpu': False
+    }
+    model = Pix2PixModel(opt)
+    model.eval()
+    image = Image.fromarray(np.array(labelmap).astype(np.uint8))
+    transform_label = get_transform(opt, method=Image.NEAREST, normalize=False)
+    # transforms.ToTensor in transform_label rescales image from [0,255] to [0.0,1.0]
+    # lets rescale it back to [0,255] to match our label ids
+    label_tensor = transform_label(image) * 255.0
+    label_tensor[label_tensor == 255] = opt['label_nc']  # 'unknown' is opt.label_nc
+    print("label_tensor:", label_tensor.shape)
+
+    # not using encoder, so creating a blank image...
+    transform_image = get_transform(opt)
+    image_tensor = transform_image(Image.new('RGB', (500, 500)))
+
+    data = {
+        'label': label_tensor.unsqueeze(0),
+        'instance': label_tensor.unsqueeze(0),
+        'image': image_tensor.unsqueeze(0)
+    }
+    generated = model(data, mode='inference')
+    print("generated_image:", generated.shape)
+
+    return generated
+
+
+def to_image(generated):
+    to_img = ToPILImage()
+    normalized_img = ((generated.reshape([3, 512, 512]) + 1) / 2.0) * 255.0
+    return to_img(normalized_img.byte().cpu())