Merge pull request #28 from open-mmlab/pytorch-1.0

Add compatibility with PyTorch 0.4.1 and 1.0

mmcv/image/io.py

@@ -54,7 +54,7 @@ def imfrombytes(content, flag='color'):
     Returns:
         ndarray: Loaded image array.
     """
-    img_np = np.fromstring(content, np.uint8)
+    img_np = np.frombuffer(content, np.uint8)
     flag = imread_flags[flag] if is_str(flag) else flag
     img = cv2.imdecode(img_np, flag)
     return img

mmcv/parallel/distributed.py

@@ -33,7 +33,10 @@ class MMDistributedDataParallel(nn.Module):
             self._dist_broadcast_coalesced(module_states,
                                            self.broadcast_bucket_size)
         if self.broadcast_buffers:
+            if torch.__version__ < '1.0':
                 buffers = [b.data for b in self.module._all_buffers()]
+            else:
+                buffers = [b.data for b in self.module.buffers()]
             if len(buffers) > 0:
                 self._dist_broadcast_coalesced(buffers,
                                                self.broadcast_bucket_size)

mmcv/runner/utils.py

@@ -5,6 +5,7 @@ from getpass import getuser
 from socket import gethostname
 
 import mmcv
+import torch
 import torch.distributed as dist
 
 
@@ -13,7 +14,11 @@ def get_host_info():
 
 
 def get_dist_info():
-    if dist._initialized:
+    if torch.__version__ < '1.0':
+        initialized = dist._initialized
+    else:
+        initialized = dist.is_initialized()
+    if initialized:
         rank = dist.get_rank()
         world_size = dist.get_world_size()
     else:

mmcv/version.py

-__version__ = '0.2.2'
+__version__ = '0.2.3'