mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-12 12:47:21 +00:00
[test] align model name with the file name. (#2045)
This commit is contained in:
@@ -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
|
||||
|
@@ -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()
|
Reference in New Issue
Block a user