Optimize examples in docstring (#766)

Signed-off-by: Jael Gu <[email protected]>

towhee/trainer/modelcard.py

@@ -65,6 +65,10 @@ class ModelCard:
         >>> # Print out model name stored in model card
         >>> model_card.model_name
         'test'
+        >>> # Save model card to "path/to/my_dir" as README.md
+        >>> model_card.save_model_card('/path/to/my_dir')
+        >>> # Save model card as "/path/to/my_dir/model_card.md"
+        >>> model_card.save_model_card('/path/to/my_dir/model_card.md')
     """
     model_name: Optional[str] = None
     model_architecture: Optional[str] = None
@@ -98,14 +102,6 @@ def save_model_card(self, save_directory_or_file):
         Args:
             save_directory_or_file (`str`):
                 file path or directory to write and save model card.
-        Return:
-            a model card written with its attributes
-
-        Example:
-            >>> from towhee.trainer.modelcard import ModelCard
-            >>> model_card = ModelCard(model_name='test', tasks='classification: num_class=10')
-            >>> model_card.save_model_card('/path/to/test/directory/')
-            Model card is saved as 'README.md' under '/path/to/test/directory/'
         """
         model_card = f"# {self.model_name}"
         model_card += f"\n\n## Model overview\n{self.model_overview}\n"

towhee/trainer/utils/layer_freezer.py

@@ -16,7 +16,24 @@ class LayerFreezer:
         >>> from towhee.models.vit.vit import VitModel
         >>> my_model = VitModel('vit_base_patch16_224')
         >>> my_freezer = LayerFreezer(my_model)
-        A LayerFreezer `my_freezer` is created for the model `my_model`.
+        >>> # Check if modules in the last layer are frozen
+        >>> my_freezer.status(-1)
+        >>> # Check if modules in the layer "head" are frozen
+        >>> my_freezer.status("head")
+        ['unfrozen', 'unfrozen']
+        >>> # Show all frozen layers
+        >>> my_freezer.show_frozen_layers()
+        ['patch_embed', 'head']
+        >>> # Freeze layers by a list of layer indexes
+        >>> my_freezer.by_idx([0, -1])
+        >>> # Freeze layers by a list of layer names
+        >>> my_freezer.by_names(['head'])
+        >>> # Freeze all layers
+        >>> my_freezer.set_all()
+        >>> # Unfreeze all layers
+        >>> my_freezer.set_all(freeze=False)
+        >>> # Freeze all except the last layer
+        >>> my_freezer.set_slice(-1)
     """
 
     def __init__(self, model):
@@ -33,17 +50,6 @@ def status(self, layer: Union[str, int]):
                 the name or index of layer.
         Return:
             A list of status ('frozen' or 'unfrozen') to indicate if modules in the layer are frozen or not.
-
-        Example:
-            >>> from towhee.trainer.utils.layer_freezer import LayerFreezer
-            >>> from towhee.models.vit.vit import VitModel
-            >>> my_model = VitModel('vit_base_patch16_224')
-            >>> my_freezer = LayerFreezer(my_model)
-            >>> # Check if modules in the last layer are frozen
-            >>> my_freezer.status(-1)
-            >>> # Check if modules in the layer "head" are frozen
-            >>> my_freezer.status("head")
-            ['unfrozen', 'unfrozen']
         """
         if isinstance(layer, str):
             if layer not in self.layer_names:
@@ -70,15 +76,6 @@ def show_frozen_layers(self):
             None.
         Return:
             A list of names of frozen layers
-
-        Example:
-            >>> from towhee.trainer.utils.layer_freezer import LayerFreezer
-            >>> from towhee.models.vit.vit import VitModel
-            >>> my_model = VitModel('vit_base_patch16_224')
-            >>> my_freezer = LayerFreezer(my_model)
-            >>> my_freezer.by_idx([0, -1])
-            >>> my_freezer.show_frozen_layers()
-            ['patch_embed', 'head']
         """
         outs = []
         for name, layer in self.model.named_children():
@@ -100,17 +97,6 @@ def by_names(self, names: list, freeze: bool = True):
                 a list of layer names
             freeze (`bool`):
                 if or not freeze layers (default: True)
-        Return:
-            None.
-
-        Example:
-            >>> from towhee.trainer.utils.layer_freezer import LayerFreezer
-            >>> from towhee.models.vit.vit import VitModel
-            >>> my_model = VitModel('vit_base_patch16_224')
-            >>> my_freezer = LayerFreezer(my_model)
-            >>> my_freezer.by_names(['head'])
-            >>> my_freezer.status('head')
-            ['frozen', 'frozen']
         """
         if not set(names).issubset(set(self.layer_names)):
             invalid_names = set(names) - set(self.layer_names)
@@ -131,17 +117,6 @@ def by_idx(self, idx: list, freeze: bool = True):
                 a list of layer indexes
             freeze (`bool`):
                 if or not freeze layers (default: True)
-        Return:
-            None.
-
-        Example:
-            >>> from towhee.trainer.utils.layer_freezer import LayerFreezer
-            >>> from towhee.models.vit.vit import VitModel
-            >>> my_model = VitModel('vit_base_patch16_224')
-            >>> my_freezer = LayerFreezer(my_model)
-            >>> my_freezer.by_names([-1])
-            >>> my_freezer.status(-1)
-            ['frozen', 'frozen']
         """
         for i in idx:
             for param in list(self.model.children())[i].parameters():
@@ -154,18 +129,6 @@ def set_all(self, freeze: bool = True):
         Args:
             freeze (`bool`):
                 if or not freeze layers (default: True)
-        Return:
-            None.
-
-        Example:
-            >>> from towhee.trainer.utils.layer_freezer import LayerFreezer
-            >>> from towhee.models.vit.vit import VitModel
-            >>> my_model = VitModel('vit_base_patch16_224')
-            >>> my_freezer = LayerFreezer(my_model)
-            >>> # Freeze all layers
-            >>> my_freezer.set_all()
-            >>> # Unfreeze all layers
-            >>> my_freezer.set_all(freeze=False)
         """
         for layer in self.model.children():
             for param in layer.parameters():
@@ -180,16 +143,6 @@ def set_slice(self, slice_num: int, freeze: bool = True):
                 number to slice the list of layers
             freeze (`bool`):
                 if or not freeze layers (default: True)
-        Return:
-            None.
-
-        Example:
-            >>> from towhee.trainer.utils.layer_freezer import LayerFreezer
-            >>> from towhee.models.vit.vit import VitModel
-            >>> my_model = VitModel('vit_base_patch16_224')
-            >>> my_freezer = LayerFreezer(my_model)
-            >>> # Freeze all except the last layer
-            >>> my_freezer.set_slice(-1)
         """
         myslice = slice(slice_num)
         slice_idx = list(range(self.layer_count))[myslice]