From f9e43d09160d0a6296324d7930d42edd76f79e96 Mon Sep 17 00:00:00 2001 From: Tal Ben-Nun Date: Fri, 2 Feb 2024 16:39:22 -0800 Subject: [PATCH] General-purpose experiment file loading API --- python/lbann/proto/serialize.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/python/lbann/proto/serialize.py b/python/lbann/proto/serialize.py index 63b3f7f60f5..37d9eca0c7f 100644 --- a/python/lbann/proto/serialize.py +++ b/python/lbann/proto/serialize.py @@ -80,3 +80,21 @@ def bin2text(infile: str, outfile: str): f.write( google.protobuf.text_format.MessageToString( message, use_index_order=True).encode()) + + +def generic_load(filename: str): + """ + Loads a .protobin or .prototext file. + """ + try: # Try binary first + message = lbann_pb2.LbannPB() + + # Read file + with open(filename, 'rb') as f: + message.ParseFromString(f.read()) + except: # Try text + with open(filename, 'rb') as f: + message = google.protobuf.text_format.Parse( + f.read(), lbann_pb2.LbannPB()) + + return message