Merge pull request #369 from denghuilu/1.2

convert dp1.2 model to dp1.3 model

Author: amcadmus

source/train/CMakeLists.txt

@@ -2,7 +2,7 @@
 
 configure_file("RunOptions.py.in" "${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py" @ONLY)
 
-file(GLOB LIB_PY main.py common.py env.py compat.py calculator.py Network.py Deep*.py Data.py DataSystem.py Model*.py Descrpt*.py Fitting.py Loss.py LearningRate.py Trainer.py TabInter.py EwaldRecp.py DataModifier.py ${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py transform.py)
+file(GLOB LIB_PY main.py common.py env.py compat.py calculator.py Network.py Deep*.py Data.py DataSystem.py Model*.py Descrpt*.py Fitting.py Loss.py LearningRate.py Trainer.py TabInter.py EwaldRecp.py DataModifier.py ${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py transform.py convert_to_13.py)
 
 file(GLOB CLS_PY  Local.py Slurm.py)
 

source/train/convert_to_13.py

@@ -0,0 +1,48 @@
+from deepmd.env import tf
+from google.protobuf import text_format
+from tensorflow.python.platform import gfile
+from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.framework import graph_util
+
+def convert_to_13(args):
+    convert_pb_to_pbtxt(args.input_model, 'frozen_model.pbtxt')
+    convert_to_dp13('frozen_model.pbtxt')
+    convert_pbtxt_to_pb('frozen_model.pbtxt', args.output_model)
+    print("the converted output model(1.3 support) is saved in %s" % args.output_model)
+
+def convert_pb_to_pbtxt(pbfile, pbtxtfile):
+    with gfile.FastGFile(pbfile, 'rb') as f:
+        graph_def = tf.GraphDef()
+        graph_def.ParseFromString(f.read())
+        tf.import_graph_def(graph_def, name='')
+        tf.train.write_graph(graph_def, './', pbtxtfile, as_text=True)
+
+def convert_pbtxt_to_pb(pbtxtfile, pbfile):
+    with tf.gfile.FastGFile(pbtxtfile, 'r') as f:
+        graph_def = tf.GraphDef()
+        file_content = f.read()
+        # Merges the human-readable string in `file_content` into `graph_def`.
+        text_format.Merge(file_content, graph_def)
+        tf.train.write_graph(graph_def, './', pbfile, as_text=False)
+
+def convert_to_dp13(file):
+    file_data = ""
+    with open(file, "r", encoding="utf-8") as f:
+        ii = 0
+        lines = f.readlines()
+        while (ii < len(lines)):
+            line = lines[ii]
+            file_data += line
+            ii+=1
+            if 'name' in line and ('DescrptSeA' in line or 'ProdForceSeA' in line or 'ProdVirialSeA' in line):
+                while not('attr' in lines[ii] and '{' in lines[ii]):
+                    file_data += lines[ii]
+                    ii+=1
+                file_data += '  attr {\n'
+                file_data += '    key: \"T\"\n'
+                file_data += '    value {\n'
+                file_data += '      type: DT_DOUBLE\n'
+                file_data += '    }\n'
+                file_data += '  }\n'
+    with open(file, "w", encoding="utf-8") as f:
+        f.write(file_data)

source/train/main.py

@@ -5,6 +5,7 @@
 from .config import config
 from .test import test
 from .transform import transform
+from .convert_to_13 import convert_to_13
 
 def main () :    
     parser = argparse.ArgumentParser(
@@ -57,6 +58,11 @@ def main () :
     parser_tst.add_argument("-d", "--detail-file", type=str, 
                             help="The file containing details of energy force and virial accuracy")
 
+    parser_transform = subparsers.add_parser('convert-to-1.3', help='convert dp-1.2 model to dp-1.3 model')
+    parser_transform.add_argument('-i', "--input-model", default = "frozen_model.pb", type=str, 
+				  help = "the input dp-1.2 model")
+    parser_transform.add_argument("-o","--output-model", default = "frozen_model_1.3.pb", type=str, 
+				  help='the converted dp-1.3 model')
     args = parser.parse_args()
 
     if args.command is None :
@@ -72,5 +78,7 @@ def main () :
         test(args)
     elif args.command == 'transform' :
         transform(args)
+    elif args.command == 'convert-to-1.3' :
+        convert_to_13(args)
     else :
         raise RuntimeError('unknown command ' + args.command)