From 29ade53c797c1d1cc1b2b977107a480f3d3f89d9 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Fri, 4 Mar 2022 14:33:55 +0800 Subject: [PATCH] Optimize examples in docstring (#766) Signed-off-by: Jael Gu --- towhee/trainer/modelcard.py | 12 ++-- towhee/trainer/utils/layer_freezer.py | 83 ++++++--------------------- 2 files changed, 22 insertions(+), 73 deletions(-) diff --git a/towhee/trainer/modelcard.py b/towhee/trainer/modelcard.py index fb760b0eca..ba696a6969 100644 --- a/towhee/trainer/modelcard.py +++ b/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" diff --git a/towhee/trainer/utils/layer_freezer.py b/towhee/trainer/utils/layer_freezer.py index ef0d21a39b..784941bee1 100644 --- a/towhee/trainer/utils/layer_freezer.py +++ b/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]