23 lines
961 B
Python
23 lines
961 B
Python
from train_prop import train_prop_model
|
|
|
|
import yaml
|
|
import argparse
|
|
|
|
# Set up argument parser
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='Training configuration of PotNet')
|
|
parser.add_argument('--config', type=str, help='Path to YAML configuration file')
|
|
parser.add_argument('--output_dir', type=str, default='output', help='Path to the output')
|
|
parser.add_argument('--data_root', type=str, default=None, help='Path to the data')
|
|
parser.add_argument('--checkpoint', type=str, default=None, help='Path to the checkpoint')
|
|
parser.add_argument('--testing', action='store_true', help='Evaluation phase')
|
|
|
|
# Parse arguments
|
|
args = parser.parse_args()
|
|
|
|
# Load data from YAML file
|
|
with open(args.config, 'r') as file:
|
|
data = yaml.safe_load(file)
|
|
|
|
data["output_dir"] = args.output_dir
|
|
train_prop_model(data, data_root=args.data_root, checkpoint=args.checkpoint, testing=args.testing) |