Write the inference and relax times to file
This commit is contained in:
parent
6e930a6ca4
commit
32dd89bb1c
|
@ -120,6 +120,7 @@ def run_model(model, batch, tag, args):
|
|||
out = model(batch)
|
||||
inference_time = time.perf_counter() - t
|
||||
logger.info(f"Inference time: {inference_time}")
|
||||
update_timings({"inference": inference_time}, os.path.join(args.output_dir, "timings.json"))
|
||||
|
||||
model.config.template.enabled = template_enabled
|
||||
|
||||
|
@ -480,7 +481,10 @@ def main(args):
|
|||
os.environ["CUDA_VISIBLE_DEVICES"] = device_no
|
||||
relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices
|
||||
logger.info(f"Relaxation time: {time.perf_counter() - t}")
|
||||
relaxation_time = time.perf_counter() - t
|
||||
|
||||
logger.info(f"Relaxation time: {relaxation_time}")
|
||||
update_timings({"relaxation": relaxation_time}, os.path.join(args.output_dir, "timings.json"))
|
||||
|
||||
# Save the relaxed PDB.
|
||||
relaxed_output_path = os.path.join(
|
||||
|
@ -500,6 +504,18 @@ def main(args):
|
|||
|
||||
logger.info(f"Model output written to {output_dict_path}...")
|
||||
|
||||
def update_timings(dict, output_file=os.path.join(os.getcwd(), "timings.json")):
|
||||
"""Write dictionary of one or more run step times to a file"""
|
||||
import json
|
||||
if os.path.exists(output_file):
|
||||
with open(output_file, "r") as f:
|
||||
timings = json.loads(f)
|
||||
else:
|
||||
timings = {}
|
||||
timings.update(dict)
|
||||
with open(output_file, "w") as f:
|
||||
timings.dumps(f)
|
||||
return output_file
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
|
|
Loading…
Reference in New Issue