add comments

dygraph/README.md

@@ -6,7 +6,7 @@ export PYTHONPATH=$PYTHONPATH:`pwd`
 
 ## 训练
 ```
-python3 train.py --model_name UNet \
+python3 train.py --model_name unet \
 --dataset OpticDiscSeg \
 --input_size 192 192 \
 --iters 10 \
@@ -17,7 +17,7 @@ python3 train.py --model_name UNet \
 
 ## 评估
 ```
-python3 val.py --model_name UNet \
+python3 val.py --model_name unet \
 --dataset OpticDiscSeg \
 --input_size 192 192 \
 --model_dir output/best_model
@@ -25,7 +25,7 @@ python3 val.py --model_name UNet \
 
 ## 预测
 ```
-python3 infer.py --model_name UNet \
+python3 infer.py --model_name unet \
 --dataset OpticDiscSeg \
 --model_dir output/best_model \
 --input_size 192 192

dygraph/models/architectures/hrnet.py

@@ -32,7 +32,23 @@ __all__ = [
 
 class HRNet(fluid.dygraph.Layer):
     """
-    HRNet：
+    HRNet：Deep High-Resolution Representation Learning for Visual Recognition
+    https://arxiv.org/pdf/1908.07919.pdf.
+
+    Args:
+        stage1_num_modules (int): number of modules for stage1. Default 1.
+        stage1_num_blocks (list): number of blocks per module for stage1. Default [4].
+        stage1_num_channels (list): number of channels per branch for stage1. Default [64].
+        stage2_num_modules (int): number of modules for stage2. Default 1.
+        stage2_num_blocks (list): number of blocks per module for stage2. Default [4, 4]
+        stage2_num_channels (list): number of channels per branch for stage2. Default [18, 36].
+        stage3_num_modules (int): number of modules for stage3. Default 4.
+        stage3_num_blocks (list): number of blocks per module for stage3. Default [4, 4, 4]
+        stage3_num_channels (list): number of channels per branch for stage3. Default [18, 36, 72].
+        stage4_num_modules (int): number of modules for stage4. Default 3.
+        stage4_num_blocks (list): number of blocks per module for stage4. Default [4, 4, 4, 4]
+        stage4_num_channels (list): number of channels per branch for stage4. Default [18, 36, 72. 144].
+        has_se (bool): whether to use Squeeze-and-Excitation module. Default False.
     """
 
     def __init__(self,

dygraph/models/fcn.py

@@ -41,6 +41,10 @@ class FCN(fluid.dygraph.Layer):
     Args:
         backbone (str): backbone name,
         num_classes (int): the unique number of target classes.
+        in_channels (int): the channels of input feature maps.
+        channels (int): channels after conv layer before the last one.
+        pretrained_model (str): the path of pretrained model.
+        ignore_index (int): the value of ground-truth mask would be ignored while computing loss or doing evaluation. Default 255.
     """
 
     def __init__(self,
@@ -49,7 +53,6 @@ class FCN(fluid.dygraph.Layer):
                  in_channels,
                  channels=None,
                  pretrained_model=None,
-                 has_se=False,
                  ignore_index=255,
                  **kwargs):
         super(FCN, self).__init__()

dygraph/models/unet.py

@@ -12,13 +12,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 import paddle.fluid as fluid
 from paddle.fluid.dygraph import Conv2D, Pool2D
-from paddle.fluid.dygraph import SyncBatchNorm as BatchNorm
+from paddle.nn import SyncBatchNorm as BatchNorm
+
+from dygraph.cvlibs import manager
+from dygraph import utils
 
 
 class UNet(fluid.dygraph.Layer):
-    def __init__(self, num_classes, ignore_index=255):
+    """
+    U-Net: Convolutional Networks for Biomedical Image Segmentation.
+    https://arxiv.org/abs/1505.04597
+
+    Args:
+        num_classes (int): the unique number of target classes.
+        pretrained_model (str): the path of pretrained model.
+        ignore_index (int): the value of ground-truth mask would be ignored while computing loss or doing evaluation. Default 255.
+    """
+
+    def __init__(self, num_classes, pretrained_model=None, ignore_index=255):
         super(UNet, self).__init__()
         self.encode = UnetEncoder()
         self.decode = UnetDecode()
@@ -26,6 +41,8 @@ class UNet(fluid.dygraph.Layer):
         self.ignore_index = ignore_index
         self.EPS = 1e-5
 
+        self.init_weight(pretrained_model)
+
     def forward(self, x, label=None):
         encode_data, short_cuts = self.encode(x)
         decode_data = self.decode(encode_data, short_cuts)
@@ -39,6 +56,20 @@ class UNet(fluid.dygraph.Layer):
             pred = fluid.layers.unsqueeze(pred, axes=[3])
             return pred, score_map
 
+    def init_weight(self, pretrained_model=None):
+        """
+        Initialize the parameters of model parts.
+        Args:
+            pretrained_model ([str], optional): the pretrained_model path of backbone. Defaults to None.
+        """
+        if pretrained_model is not None:
+            if os.path.exists(pretrained_model):
+                utils.load_pretrained_model(self.backbone, pretrained_model)
+                utils.load_pretrained_model(self, pretrained_model)
+            else:
+                raise Exception('Pretrained model is not found: {}'.format(
+                    pretrained_model))
+
     def _get_loss(self, logit, label):
         logit = fluid.layers.transpose(logit, [0, 2, 3, 1])
         label = fluid.layers.transpose(label, [0, 2, 3, 1])
@@ -108,14 +139,14 @@ class DoubleConv(fluid.dygraph.Layer):
             filter_size=3,
             stride=1,
             padding=1)
-        self.bn0 = BatchNorm(num_channels=num_filters)
+        self.bn0 = BatchNorm(num_filters)
         self.conv1 = Conv2D(
             num_channels=num_filters,
             num_filters=num_filters,
             filter_size=3,
             stride=1,
             padding=1)
-        self.bn1 = BatchNorm(num_channels=num_filters)
+        self.bn1 = BatchNorm(num_filters)
 
     def forward(self, x):
         x = self.conv0(x)
@@ -166,3 +197,8 @@ class GetLogit(fluid.dygraph.Layer):
     def forward(self, x):
         x = self.conv(x)
         return x
+
+
+@manager.MODELS.add_component
+def unet(*args, **kwargs):
+    return UNet(*args, **kwargs)