[Feature Request] Support different backends for DeePMD-kit

[njzjz]

Summary

Support PyTorch backend for DeePMD-kit.

Detailed Description

Although I have submitted deepmodeling/deepmd-kit#3191 to make most of the things the same, the suffixes of the model and checkpoint files are different by design. These suffixes are hardcoded in the code.

dpgen/dpgen/generator/run.py

Lines 804 to 812 in 7392650

	backward_files = ["frozen_model.pb", "lcurve.out", "train.log"]
	backward_files += [
	"model.ckpt.meta",
	"model.ckpt.index",
	"model.ckpt.data-00000-of-00001",
	"checkpoint",
	]
	if jdata.get("dp_compress", False):
	backward_files.append("frozen_model_compressed.pb")

dpgen/dpgen/generator/run.py

Lines 884 to 889 in 7392650

	if not jdata.get("dp_compress", False):
	model_name = "frozen_model.pb"
	else:
	model_name = "frozen_model_compressed.pb"
	task_file = os.path.join(train_task_fmt % ii, model_name)
	ofile = os.path.join(work_path, "graph.%03d.pb" % ii)

dpgen/dpgen/generator/run.py

Line 1938 in 7392650

all_models = glob.glob(os.path.join(work_path, "graph*pb"))

dpgen/dpgen/simplify/simplify.py

Line 188 in 7392650

models = glob.glob(os.path.join(train_path, "graph*pb"))

Besides, glob doesn't work for backward_files - in dpdispatcher, some contexts like Bohrium don't support it.
So, may we consider adding a parameter to switch?

Further Information, Files, and Links

No response

[njzjz]

This is blocked by deepmodeling/deepmd-kit#3475.

[thangckt]

hi @njzjz

Did you solved this problem?
I made some changes to have option for choosing between backends TF and PT, in this file

If you find it helpful, I will make a PR. Otherwise, just omit

[njzjz]

Did you solved this problem?

I don't have plans for this issue.

Contribution is welcome.