coderabbitai[bot]
commented
1 month ago Walkthrough
The changes introduce a new command-line argument --list-model-branch to list model branches in a multitask pretrained model, supported by the PyTorch backend. The updates include handling this flag in the training command, adding a display function for model information, and enhancing test cases for single-task and multi-task models. Additionally, a utility function to run DP directly from the entry point has been added to improve testing performance.
Changes
| File/Directory | Summary |
|---|---|
deepmd/main.py |
Added --list-model-branch argument to parser_train for listing model branches in multitask pretrained models. |
deepmd/pt/entrypoints/main.py |
Added logic to handle --list-model-branch flag in the train command and a show function to display model info. |
doc/train/finetuning.md |
Updated documentation to include the new --list-model-branch command for checking model branches. |
source/tests/pt/common.py |
Introduced run_dp function to run DP directly from the entry point, improving performance by avoiding subprocess use. |
source/tests/pt/test_dp_show.py |
Added test cases for single-task and multi-task models, including setup, training, and model information display. |
Sequence Diagram(s) (Beta)
sequenceDiagram
participant User
participant CLI
participant Main
participant EntryPoints
User->>CLI: Execute train command with --list-model-branch
CLI->>Main: Parse arguments
Main->>EntryPoints: Call train function with flags
EntryPoints-->>EntryPoints: Check --list-model-branch flag
alt Flag is set
EntryPoints->>EntryPoints: Process pretrained model for multitask mode
EntryPoints->>EntryPoints: Extract and display model branches
else Flag is not set
EntryPoints->>EntryPoints: Proceed with regular training process
end
EntryPoints-->>Main: Return control
Main-->>CLI: Finish execution
CLI-->>User: Display resultsThank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?
codecov[bot]
Chengqian-Zhang
wanghan-iapcm
Solve #3742
Situation one（The right way to use it）：
dp --pt show multitask_model.pt model-branch type-map descriptor fitting-net[2024-05-22 10:38:16,678] DEEPMD INFO This is a multitask model[2024-05-22 10:38:16,678] DEEPMD INFO Available model branches are ['MPtraj_v026_01-mix-Utype', 'MPtraj_v026_02-mix-Utype', 'MPtraj_v026_03-mix-Utype', 'MPtraj_v026_04-mix-Utype', 'MPtraj_v026_05-mix-Utype', 'MPtraj_v026_06-mix-Utype', 'MPtraj_v026_07-mix-Utype', 'MPtraj_v026_08-mix-Utype', 'MPtraj_v026_09-mix-Utype', 'MPtraj_v026_10-mix-Utype', 'MPtraj_v026_11-mix-Utype'][2024-05-22 10:38:16,679] DEEPMD INFO The type_map of branch MPtraj_v026_01-mix-Utype is ['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og', 'Co_U', 'Cr_U', 'Fe_U', 'Mn_U', 'Mo_U', 'Ni_U', 'V_U', 'W_U'](skip other branches' output)[2024-05-22 10:38:16,679] DEEPMD INFO The descriptor parameter of branch MPtraj_v026_04-mix-Utype is {'type': 'dpa2', 'repinit': {'tebd_dim': 256, 'rcut': 9.0, 'rcut_smth': 8.0, 'nsel': 120, 'neuron': [25, 50, 100], 'axis_neuron': 12, 'activation_function': 'tanh'}, 'repformer': {'rcut': 4.0, 'rcut_smth': 3.5, 'nsel': 40, 'nlayers': 12, 'g1_dim': 128, 'g2_dim': 32, 'attn2_hidden': 32, 'attn2_nhead': 4, 'attn1_hidden': 128, 'attn1_nhead': 4, 'axis_neuron': 4, 'activation_function': 'tanh', 'update_h2': False, 'update_g1_has_conv': True, 'update_g1_has_grrg': True, 'update_g1_has_drrd': True, 'update_g1_has_attn': True, 'update_g2_has_g1g1': False, 'update_g2_has_attn': True, 'update_style': 'res_residual', 'update_residual': 0.01, 'update_residual_init': 'norm', 'attn2_has_gate': True}, 'add_tebd_to_repinit_out': False}(skip other branches' output)[2024-05-22 10:38:16,679] DEEPMD INFO The fitting_net parameter of branch MPtraj_v026_01-mix-Utype is {'neuron': [240, 240, 240], 'activation_function': 'tanh', 'resnet_dt': True, 'seed': 1, '_comment': " that's all"}(skip other branches' output)Situation two (
singletask_model.ptis not a multi-task pre-trained model)dp --pt show singletask_model.pt model-branch type-map descriptor fitting-net[2024-05-22 10:43:11,642] DEEPMD INFO This is a singletask modelRuntimeError: The 'model-branch' option requires a multitask model. The provided model does not meet this criterion.Situation three(using tf backend)
dp show multitask_model.pt model-branchRuntimeError: unknown command list-model-branchFrozen model file with a .pth extension are used in the same way as checkpoint file with a .pt extension.
dp --pt show frozen_model.pth type-map descriptor fitting-net[2024-05-22 10:46:26,365] DEEPMD INFO This is a singletask model[2024-05-22 10:46:26,365] DEEPMD INFO The type_map is ['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og', 'Co_U', 'Cr_U', 'Fe_U', 'Mn_U', 'Mo_U', 'Ni_U', 'V_U', 'W_U'][2024-05-22 10:46:26,365] DEEPMD INFO The descriptor parameter is {'type': 'dpa2', 'repinit': {'tebd_dim': 256, 'rcut': 9.0, 'rcut_smth': 8.0, 'nsel': 120, 'neuron': [25, 50, 100], 'axis_neuron': 12, 'activation_function': 'tanh'}, 'repformer': {'rcut': 4.0, 'rcut_smth': 3.5, 'nsel': 40, 'nlayers': 12, 'g1_dim': 128, 'g2_dim': 32, 'attn2_hidden': 32, 'attn2_nhead': 4, 'attn1_hidden': 128, 'attn1_nhead': 4, 'axis_neuron': 4, 'activation_function': 'tanh', 'update_h2': False, 'update_g1_has_conv': True, 'update_g1_has_grrg': True, 'update_g1_has_drrd': True, 'update_g1_has_attn': True, 'update_g2_has_g1g1': False, 'update_g2_has_attn': True, 'update_style': 'res_residual', 'update_residual': 0.01, 'update_residual_init': 'norm', 'attn2_has_gate': True}, 'add_tebd_to_repinit_out': False}[2024-05-22 10:46:26,365] DEEPMD INFO The fitting_net parameter is {'neuron': [240, 240, 240], 'activation_function': 'tanh', 'resnet_dt': True, 'seed': 1, '_comment': " that's all"}Summary by CodeRabbit
New Features
--list-model-branchcommand-line argument for listing model branches of a multitask pretrained model.Documentation
--list-model-branchfor checking available model branches in a multitask pre-trained model.Tests