1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
|
from .criterions import cel
from .ddp_models import basic_ddp_model
from .hook_states import BasicHookState
from .hooks import allreduce_hook, hybrid_hook, rpc_hook, sparse_rpc_hook
from .iteration_steps import basic_iteration_step
from .preprocess_data import preprocess_dummy_data
from .trainer import DdpTrainer
criterion_map = {
"cel": cel
}
ddp_hook_map = {
"allreduce_hook": allreduce_hook,
"hybrid_hook": hybrid_hook,
"rpc_hook": rpc_hook,
"sparse_rpc_hook": sparse_rpc_hook
}
ddp_model_map = {
"basic_ddp_model": basic_ddp_model
}
iteration_step_map = {
"basic_iteration_step": basic_iteration_step
}
preprocess_data_map = {
"preprocess_dummy_data": preprocess_dummy_data
}
hook_state_map = {
"BasicHookState": BasicHookState
}
trainer_map = {
"DdpTrainer": DdpTrainer
}
|
