update readme

Author: yi.he

README.md

@@ -14,9 +14,9 @@
 
 * ``alignshift``
   the core implementation of AlignShift convolution and TSM convolution, including the operators, models, and 2D-to-3D/AlignShift/TSM model converters. 
-  * ``operators``: include AlignShiftConv, TSMConv.
-  * ``converters.py``: include converters which convert 2D models to 3dConv/AlignShiftConv/TSMConv counterparts.
-  * ``models``: Native AlignShift/TSM models. 
+  * ``operators``: include AlignShiftConv, TSMConv, A3DConv.
+  * ``converters.py``: include converters which convert 2D models to 3DConv/AlignShiftConv/TSMConv/A3DConv counterparts.
+  * ``models``: Native AlignShift/TSM/A3DConv models. 
 * ``deeplesion`` 
   the experiment code is base on [mmdetection](https://github.com/open-mmlab/mmdetection)
 ,this directory consists of compounents used in mmdetection.
@@ -51,10 +51,10 @@ thickness = torch.rand(batch_size, 1)
 out = m(x, thickness)
 ```
 
-## Usage of AlignShiftConv/TSMConv operators
+## Usage of AlignShiftConv/TSMConv/A3DConv operators
 
 ```python
-from nn.operators import AlignShiftConv, TSMConv
+from nn.operators import AlignShiftConv, TSMConv, A3DConv
 x = torch.rand(batch_size, 3, D, H, W)
 thickness = torch.rand(batch_size, 1)
 # AlignShiftConv to process 3D volumnes
@@ -63,6 +63,9 @@ out = conv(x, thickness)
 # TSMConv to process 3D volumnes
 conv = TSMConv(in_channels=3, out_channels=10, kernel_size=3, padding=1, n_fold=8, tsm=True)
 out = conv(x)
+# A3DConv to process 3D volumnes
+conv = A3DConv(in_channels=3, out_channels=10, kernel_size=3, padding=1, dimension=3)
+out = conv(x)
 ```
 
 ## Usage of native  AlignShiftConv/TSMConv models
@@ -109,6 +112,11 @@ output_3d = net(input_3d, thickness)
   ```bash
   ./deeplesion/train_dist.sh ./deeplesion/mconfig/densenet_tsm.py 2
   ```
+    * Train A3DConv models 
+  ```bash
+  ./deeplesion/train_dist.sh ./deeplesion/mconfig/densenet_a3d.py 2
+  ```
+
  * Evaluation 
    ```bash
    ./deeplesion/eval.sh ${mmdetection script} ${checkpoint path}

nn/operators/__init__.py

@@ -1,3 +1,3 @@
 from .alignshiftconv import AlignShiftConv
 from .tsmconv import TSMConv
-from .a3d import A3DConv
+from .a3dconv import A3DConv

nn/operators/a3d.py renamed to nn/operators/a3dconv.py

@@ -10,18 +10,14 @@
 class A3DConv(_ConvNd):
     '''
     Args: 
-        n_fold (int): Divisor of channels to shift.
-        alignshift(bool): if apply alignshift operation before conv
+        dimension (int): number of dimensions to shift.
         inplace (bool): if Enable inplace operation.
-        ref_thickness (float): Reference z axis spacing Default: 0.2mm.
-        shift_padding_zero(bool): f padding zeros to side fold before shift channels 
+        enable_shift(bool): if enable shift 
     '''
     def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                  padding=0, dilation=1, groups=1,
                  bias=True, padding_mode='zeros', 
-                 dimension=3,
-                 enable_shift=True, inplace=False,
-                 shift_padding_zero=False):
+                 dimension=3, enable_shift=True, inplace=False):
 
         kernel_size = _pair_same(kernel_size)
         stride = as_triple(stride)
@@ -33,8 +29,6 @@ def __init__(self, in_channels, out_channels, kernel_size, stride=1,
 
         self.enable_shift = enable_shift
         self.inplace = inplace
-        self.shift_padding_zero = shift_padding_zero
-        # if self.enable_shift:
         self._init_adaptive_weights(in_channels, dimension)
 
     def _init_adaptive_weights(self, in_channels, dimension):