add comments

faye0078 · Oct 26, 2022 · 9611bad · 9611bad
1 parent b6ce143
commit 9611bad
Showing 1 changed file with 15 additions and 15 deletions.
diff --git a/model/edge.py b/model/edge.py
@@ -5,14 +5,14 @@
 class BasicEdge(nn.Module):
 
     def __init__(self, input_channels, output_channels):
-        """_summary_
+        """Basic Edge in supernet
 
         Args:
-            input_channels (_type_): _description_
-            output_channels (_type_): _description_
+            input_channels (int): input channel of edge
+            output_channels (int): output channel of edge
 
         Raises:
-            ValueError: _description_
+            ValueError: input_channels must be equal to output_channels
         """
         if input_channels != output_channels:
             raise ValueError(f'input_channels != output_channels: {input_channels} != {output_channels}, basic edge only support same input and output channels')
@@ -44,11 +44,11 @@ def __init__(self, input_channels, output_channels):
         """_summary_
 
         Args:
-            input_channels (_type_): _description_
-            output_channels (_type_): _description_
+            input_channels (int): input channel of edge
+            output_channels (int): output channel of edge
 
         Raises:
-            ValueError: _description_
+            ValueError: input_channels must be equal to output_channels
         """
         if input_channels != output_channels:
             raise ValueError(f'input_channels != output_channels: {input_channels} != {output_channels}, keep edge only support same input and output channels')
@@ -73,12 +73,12 @@ def __init__(self, input_channels, output_channels, down_num):
         """_summary_
 
         Args:
-            input_channels (_type_): _description_
-            output_channels (_type_): _description_
-            down_num (_type_): _description_
+            input_channels (int): input channel of edge
+            output_channels (int): output channel of edge
+            down_num (int): downsample times
 
         Raises:
-            ValueError: _description_
+            ValueError: input_channels must be less than output_channels
         """        
         if input_channels >= output_channels:
             raise ValueError(f'input_channels >= output_channels: {input_channels} >= {output_channels}, downsample node only support input_channels < output_channels')
@@ -103,12 +103,12 @@ def __init__(self, input_channels, output_channels, up_num):
         """_summary_
 
         Args:
-            input_channels (_type_): _description_
-            output_channels (_type_): _description_
-            up_num (_type_): _description_
+            input_channels (int): input channel of edge
+            output_channels (int): output channel of edge
+            up_num (int): upsample times 
 
         Raises:
-            ValueError: _description_
+            ValueError: input_channels must be greater than output_channels
         """        
         if input_channels <= output_channels:
             raise ValueError(f'input_channels <= output_channels: {input_channels} <= {output_channels}, upsample node only support input_channels > output_channels')