v0.4.2

setup.py

@@ -18,22 +18,7 @@ import setuptools.command.install
 
 cwd = os.path.dirname(os.path.abspath(__file__))
 
-# run test scrip after installation
-class install(setuptools.command.install.install):
-    def run(self):
-        self.create_version_file()
-        setuptools.command.install.install.run(self)
-        #subprocess.check_call("python tests/unit_test.py".split())
-    @staticmethod
-    def create_version_file():
-        global version, cwd
-        print('-- Building version ' + version)
-        version_path = os.path.join(cwd, 'encoding', 'version.py')
-        with open(version_path, 'w') as f:
-            f.write('"""This is encoding version file."""\n')
-            f.write("__version__ = '{}'\n".format(version))
-
-version = '0.4.0'
+version = '0.4.2'
 try:
     sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], 
         cwd=cwd).decode('ascii').strip()
@@ -41,6 +26,27 @@ try:
 except Exception:
     pass
 
+def create_version_file():
+    global version, cwd
+    print('-- Building version ' + version)
+    version_path = os.path.join(cwd, 'encoding', 'version.py')
+    with open(version_path, 'w') as f:
+        f.write('"""This is encoding version file."""\n')
+        f.write("__version__ = '{}'\n".format(version))
+
+# run test scrip after installation
+class install(setuptools.command.install.install):
+    def run(self):
+        create_version_file()
+        setuptools.command.install.install.run(self)
+        #subprocess.check_call("python tests/unit_test.py".split())
+
+class develop(setuptools.command.develop.develop):
+    def run(self):
+        create_version_file()
+        setuptools.command.develop.develop.run(self)
+        #subprocess.check_call("python tests/unit_test.py".split())
+
 try:
     import pypandoc
     readme = pypandoc.convert('README.md', 'rst')
@@ -51,10 +57,20 @@ requirements = [
     'numpy',
     'tqdm',
     'nose',
-    'torch>=0.3.1',
+    'torch>=0.5.0',
     'cffi>=1.0.0',
 ]
 
+requirements = [
+    'numpy',
+    'tqdm',
+    'nose',
+    'torch>=0.4.0',
+    'Pillow',
+    'scipy',
+    'requests',
+]
+
 setup(
     name="encoding",
     version=version,
@@ -67,17 +83,14 @@ setup(
     install_requires=requirements,
     packages=find_packages(exclude=["tests", "experiments"]),
     package_data={ 'encoding': [
-        'lib/*.so*', 'lib/*.dylib*',
-        '_ext/encoding_lib/*.so', '_ext/encoding_lib/*.dylib',
-        'kernel/*.h', 'kernel/generic/*h',
-        'src/*.h',
+        'lib/cpu/*.h',
+        'lib/cpu/*.cpp',
+        'lib/gpu/*.h',
+        'lib/gpu/*.cpp',
+        'lib/gpu/*.cu',
     ]},
-    ext_package="",
-    # Extensions to compile.
-    cffi_modules=[
-        os.path.join(cwd, "build.py:ffi")
-    ],
     cmdclass={
         'install': install,
+        'develop': develop,
     },
 )

tests/unit_test.py → tests/unit_test/test_module.py

@@ -15,12 +15,12 @@ from torch.autograd import Variable, gradcheck
 
 import encoding
 
-EPS = 1e-6
-ATOL = 1e-4
+EPS = 1e-3
+ATOL = 1e-3
 
 def _assert_tensor_close(a, b, atol=ATOL, rtol=EPS):
     npa, npb = a.cpu().numpy(), b.cpu().numpy()
-    assert np.allclose(npa, npb, atol=atol), \
+    assert np.allclose(npa, npb, rtol=rtol, atol=atol), \
         'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(
             a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max())
 
@@ -79,95 +79,92 @@ def test_all_reduce():
     for i in range(1, ngpu):
         _assert_tensor_close(Y[i].data, Y[0].data)
     input = (1, *X)
-    #test = gradcheck(encoding.parallel.allreduce, input, eps=EPS, atol=ATOL)
     test = gradcheck(encoding.parallel.allreduce, input, eps=EPS, atol=ATOL)
+    print('Testing allreduce(): {}'.format(test))
 
-def _test_syncbn(train_mode=True):
+
+def test_syncbn():
+    train_mode=True
     # generate input
     B,C,H,W = 8,3,4,5
     X = Variable(torch.cuda.DoubleTensor(B,C,H,W).uniform_(-0.5,0.5), 
                  requires_grad=True)
     input = (X,)
     # SyncBN using DataParallel
-    layer = encoding.nn.SyncBatchNorm2d(C)
+    layer = encoding.nn.BatchNorm2d(C)
     model = torch.nn.DataParallel(layer).double().cuda()
+    encoding.parallel.patch_replication_callback(model)
     layer.train(train_mode)
     # grad check
     test = gradcheck(model, input, eps=EPS, atol=ATOL)
-    print('Testing SyncBatchNorm2d(): {}'.format(test))
+    print('Testing BatchNorm2d(): {}'.format(test))
 
 
-def _test_syncbn_func(train_mode=True):
+def test_syncbn_func():
     # generate input
     B, C, H = 2, 3, 4
     X = Variable(torch.cuda.DoubleTensor(B,C,H).uniform_(-0.5, 0.5), 
         requires_grad=True)
-    xsum = Variable(torch.ones(C).double().cuda(), requires_grad=True)
-    xsqsum = Variable(torch.ones(C).double().cuda(), requires_grad=True)
-    gamma = Variable(torch.ones(C).double().cuda(), requires_grad=True)
-    beta = Variable(torch.zeros(C).double().cuda(), requires_grad=True)
-    gamma.requires_grad=True
-    beta.requires_grad=True
-    runningVar = torch.ones(C).double().cuda()
-    runningMean = torch.zeros(C).double().cuda()
+    gamma = Variable(torch.cuda.DoubleTensor(C).uniform_(-0.5, 0.5), requires_grad=True)
+    beta = Variable(torch.cuda.DoubleTensor(C).uniform_(-0.5, 0.5), requires_grad=True)
+    mean = Variable(torch.cuda.DoubleTensor(C).uniform_(-0.5, 0.5), requires_grad=True)
+    std = Variable(torch.cuda.DoubleTensor(C).uniform_(-0.5, 0.5), requires_grad=True)
     N = B * H
-    inputs = (X, xsum, xsqsum, gamma, beta, runningMean, runningVar, N, 0.1, 1e-5, train_mode)
+    inputs = (X, mean, std, gamma, beta)
     # grad check
-    test = gradcheck(encoding.functions.batchnorm.apply, inputs, eps=EPS, atol=ATOL)
+    test = gradcheck(encoding.functions.batchnormtrain, inputs, eps=EPS, atol=ATOL)
     print('Testing batchnorm(): {}'.format(test))
 
 
-def _checkBatchNormResult(bn1, bn2, input, is_train, cuda=False):
-    def _find_bn(module):
-        for m in module.modules():
-            if isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d,
-                              encoding.nn.SyncBatchNorm1d, encoding.nn.SyncBatchNorm2d)):
-                return m
-    def _syncParameters(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)
-            bn2.running_mean.copy_(bn1.running_mean)
-            bn2.running_var.copy_(bn1.running_var)
-
-    bn1.train(mode=is_train)
-    bn2.train(mode=is_train)
-
-    if cuda:
-        input = input.cuda()
-    # using the same values for gamma and beta
-    _syncParameters(_find_bn(bn1), _find_bn(bn2))
-
-    input1 = Variable(input.clone(), requires_grad=True)
-    output1 = bn1(input1)
-    input2 = Variable(input.clone(), requires_grad=True)
-    output2 = bn2(input2)
-
-    _assert_tensor_close(input1.data, input2.data)
-    _assert_tensor_close(output1.data, output2.data)
-    if not is_train:
-        return
-    (output1 ** 2).sum().backward()
-    (output2 ** 2).sum().backward()
-    _assert_tensor_close(input1.grad.data, input2.grad.data)
-    _assert_tensor_close(_find_bn(bn1).running_mean, _find_bn(bn2).running_mean)
-    _assert_tensor_close(_find_bn(bn1).running_var, _find_bn(bn2).running_var)
+def testSyncBN():
+    def _checkBatchNormResult(bn1, bn2, input, is_train, cuda=False):
+        def _find_bn(module):
+            for m in module.modules():
+                if isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d,
+                                  encoding.nn.BatchNorm1d, encoding.nn.BatchNorm2d)):
+                    return m
+        def _syncParameters(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)
+                bn2.running_mean.copy_(bn1.running_mean)
+                bn2.running_var.copy_(bn1.running_var)
+
+        bn1.train(mode=is_train)
+        bn2.train(mode=is_train)
+
+        if cuda:
+            input = input.cuda()
+        # using the same values for gamma and beta
+        _syncParameters(_find_bn(bn1), _find_bn(bn2))
+
+        input1 = Variable(input.clone().detach(), requires_grad=True)
+        input2 = Variable(input.clone().detach(), requires_grad=True)
+        output1 = bn1(input1)
+        output2 = bn2(input2)
+        # assert forwarding
+        _assert_tensor_close(input1.data, input2.data)
+        _assert_tensor_close(output1.data, output2.data)
+        if not is_train:
+            return
+        (output1 ** 2).sum().backward()
+        (output2 ** 2).sum().backward()
+        _assert_tensor_close(input1.grad.data, input2.grad.data)
+        _assert_tensor_close(_find_bn(bn1).running_mean, _find_bn(bn2).running_mean)
+        _assert_tensor_close(_find_bn(bn1).running_var, _find_bn(bn2).running_var)
 
 
-def testSyncBN():
     bn = torch.nn.BatchNorm2d(10).cuda().double()
-    sync_bn = encoding.nn.SyncBatchNorm2d(10).double()
+    sync_bn = encoding.nn.BatchNorm2d(10).double()
     sync_bn = torch.nn.DataParallel(sync_bn).cuda()
+    encoding.parallel.patch_replication_callback(sync_bn)
     # check with unsync version
     for i in range(10):
+        print(i)
         _checkBatchNormResult(bn, sync_bn, torch.rand(16, 10, 16, 16).double(), True, cuda=True)
         _checkBatchNormResult(bn, sync_bn, torch.rand(16, 10, 16, 16).double(), False, cuda=True)
-    # gradcheck
-    _test_syncbn_func(True)
-    _test_syncbn(True)
-
 
 if __name__ == '__main__':
     import nose