()
| 256 | |
| 257 | |
| 258 | def main(): |
| 259 | |
| 260 | #argparse |
| 261 | parser = argparse.ArgumentParser(description='Train sdf and color') |
| 262 | parser.add_argument('--dataset_path', required=True, help='Path to the difflocks dataset to train on') |
| 263 | parser.add_argument('--exp_info', default="", help='Experiment info string useful for distinguishing one experiment for another') |
| 264 | parser.add_argument('--device', default="cuda") |
| 265 | args = parser.parse_args() |
| 266 | |
| 267 | #get the output path which will be at the root of the package |
| 268 | hair_forge_root=os.path.dirname(os.path.abspath(__file__)) |
| 269 | output_training_path=os.path.join(hair_forge_root, "out_training") |
| 270 | os.makedirs(output_training_path, exist_ok=True) |
| 271 | |
| 272 | |
| 273 | experiment_name="hair_forge" |
| 274 | if args.exp_info: |
| 275 | experiment_name+="_"+args.exp_info |
| 276 | |
| 277 | |
| 278 | loader_train, loader_test= create_dataloaders(args.dataset_path) |
| 279 | |
| 280 | |
| 281 | hyperparams=HyperParamsStrandVAE() |
| 282 | train(args, hyperparams, loader_train, loader_test, experiment_name, output_training_path) |
| 283 | |
| 284 | #finished training |
| 285 | return |
| 286 | |
| 287 | |
| 288 | if __name__ == '__main__': |
no test coverage detected