[test] align model name with the file name. (#2045)

This commit is contained in:
Jiarui Fang
2022-11-30 15:45:26 +08:00
committed by GitHub
parent 31c644027b
commit 1e885329f4
10 changed files with 225 additions and 219 deletions

View File

@@ -1,2 +1,11 @@
from . import bert, gpt, inline_op_model, nested_model, no_leaf_module, repeated_computed_layer, resnet, simple_net
from . import (
bert,
gpt2,
hanging_param_model,
inline_op_model,
nested_model,
repeated_computed_layer,
resnet,
simple_net,
)
from .utils import run_fwd_bwd

View File

@@ -8,9 +8,10 @@ from .registry import non_distributed_component_funcs
from .utils.dummy_data_generator import DummyDataGenerator
class NoLeafModule(CheckpointModule):
class HangingParamModule(CheckpointModule):
"""
In this no-leaf module, it has subordinate nn.modules and a nn.Parameter.
Hanging Parameter: a parameter dose not belong to a leaf Module.
It has subordinate nn.modules and a nn.Parameter.
"""
def __init__(self, checkpoint=False) -> None:
@@ -34,11 +35,11 @@ class DummyDataLoader(DummyDataGenerator):
return data, label
@non_distributed_component_funcs.register(name='no_leaf_module')
@non_distributed_component_funcs.register(name='hanging_param_model')
def get_training_components():
def model_builder(checkpoint=False):
return NoLeafModule(checkpoint)
return HangingParamModule(checkpoint)
trainloader = DummyDataLoader()
testloader = DummyDataLoader()