Spaces:
Sleeping
Sleeping
| import json | |
| import sys, os | |
| # import torch | |
| # torch.multiprocessing.set_start_method('spawn') | |
| # torch.multiprocessing.set_sharing_strategy('file_system') | |
| from pytorch_lightning import seed_everything | |
| from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, StochasticWeightAveraging | |
| import pytorch_lightning as pl | |
| from pytorch_lightning.plugins import DDPPlugin | |
| from pricePrediction import config | |
| from pricePrediction.ArgParser_base import MyArgParser | |
| from pricePrediction.dataManager.dataManager import GraphPriceDatamodule | |
| from pricePrediction.evaluation.evaluation import PlotFigsTensorboardCallback, InferencePlotter | |
| from pricePrediction.nets.netsGraph import PricePredictorModule | |
| def parse_args(): | |
| parser = MyArgParser(prog=None, usage=None, description=None, ) | |
| #### Program args #### | |
| parser.add_argument("-e", "--n_epochs", type=int, help="Number of epochs to use", default=1000) | |
| parser.add_argument("-m", "--msg", type=str, help="A message describing the run", required=True) | |
| parser.add_argument("-r", "--restore", type=str, help="Directory to trained checkpoint to continue training", default=None) | |
| parser.add_argument("-c", "--config", type=str, help="Path to json file with the config arguments to create/update the " | |
| "model") | |
| parser.add_argument("-o", "--default_root_dir", type=str, help="The directory where to save " | |
| "checkpoints and logs", default=None) | |
| parser.add_argument("-g", "--gpus", type=int, help="Number of gpus to use", default=1) | |
| parser.add_argument("-n", "--num_nodes", type=int, help="Number of nodes to use", default=None) | |
| parser.add_argument("-s", "--random_seed", type=int, help="Random seed", default=121) | |
| parser.add_argument( "--limit_train_batches", type=float, help="Train only with a fraction of the tranining set", default=None) | |
| parser.add_argument("--swa", action="store_true", default=False, | |
| help="Perform Stochastic Weight Averaging") | |
| #### Model and data args #### | |
| group = parser.add_argument_group(title="data") | |
| GraphPriceDatamodule.addParamsToArgParse(group) | |
| group = parser.add_argument_group(title="model") | |
| PricePredictorModule.addParamsToArgParse(group) | |
| cmd_args = parser.parse_args() | |
| del cmd_args['positional arguments'] | |
| cmd_args['main'] = cmd_args['optional arguments'] | |
| del cmd_args['optional arguments'] | |
| return cmd_args | |
| if __name__ == "__main__": | |
| print( " ".join(sys.argv)) | |
| cmd_args = parse_args() | |
| if cmd_args.get("config"): | |
| print("Loading configuration file:", cmd_args.get("config")) | |
| with open(cmd_args.get("config")) as f: | |
| args = json.load(f) | |
| for groupName in cmd_args: | |
| args[groupName].update(cmd_args[groupName]) | |
| else: | |
| args = cmd_args | |
| print(args) | |
| seed_everything(args["main"]["random_seed"]) | |
| if args['main']["restore"]: | |
| prev_run_dir = os.path.expanduser(args['main']["restore"]) | |
| # hparams = yaml.load( os.path.join(prev_run_dir, "hparams.yaml")) | |
| checkpointsDir = os.path.join(prev_run_dir, "checkpoints") | |
| most_recent_checkpoint_fname = max([os.path.join(checkpointsDir, basename) for basename in os.listdir(checkpointsDir)], | |
| key=os.path.getctime) | |
| pl_model = PricePredictorModule.load_from_checkpoint(most_recent_checkpoint_fname, **args["model"]) | |
| data_args = pl_model.hparams.get("data_hparams") | |
| data_args.update(args["data"]) | |
| dataModule = GraphPriceDatamodule(**data_args) | |
| print("Previous model loaded!!") | |
| else: | |
| dataModule = GraphPriceDatamodule(**args["data"]) | |
| dataModule.prepare_data() | |
| nodes_n_features, edges_n_features = dataModule.dims | |
| nodes_degree = dataModule.get_nodes_degree() | |
| args["model"].update( dict( | |
| nodes_n_features=nodes_n_features, edges_n_features=edges_n_features, | |
| deg=nodes_degree, data_hparams= args["data"], logs_only_in_epoch= args["main"].get("gpus", 0) > 1) | |
| ) | |
| pl_model = PricePredictorModule( **args["model"] ) | |
| print( pl_model.hparams ) | |
| callbacks = [ | |
| EarlyStopping(monitor='val_loss', patience=60), | |
| ModelCheckpoint(monitor='val_loss', verbose=True), | |
| PlotFigsTensorboardCallback(frequency=5, save_csv=True) | |
| ] | |
| if args["main"].get("swa", False): | |
| callbacks += [ | |
| StochasticWeightAveraging(annealing_epochs=config.COSINE_LR_SCHEDULE_N_EPOCHS, swa_lrs= pl_model.hparams.lr) | |
| ] | |
| default_root_dir = args['main'].get('default_root_dir', os.getcwd()) | |
| trainer_args = dict(gpus=1, max_epochs=args['main']["n_epochs"], progress_bar_refresh_rate=20, | |
| default_root_dir=default_root_dir, auto_lr_find=True, | |
| callbacks = callbacks) | |
| if args["main"]["limit_train_batches"]: | |
| trainer_args["limit_train_batches"] = args["main"]["limit_train_batches"] | |
| ngpus = args['main'].get("gpus", -1) | |
| if ngpus > 1: | |
| trainer_args["gpus"] = ngpus | |
| trainer_args["accelerator"] = "ddp" | |
| trainer_args["plugins"] = DDPPlugin(find_unused_parameters=False), | |
| os.environ["NCCL_NSOCKS_PERTHREAD"]="4" | |
| os.environ["NCCL_SOCKET_NTHREADS"] = "2" | |
| elif ngpus ==0: | |
| trainer_args["gpus"] = 0 | |
| trainer_args["accelerator"] = "ddp_cpu" | |
| if args['main'].get("num_nodes", None): | |
| trainer_args["num_nodes"] = int(args['main'].get("num_nodes")) | |
| trainer_args["accelerator"] = "ddp" | |
| raise NotImplementedError() | |
| # import torch.distributed as dist | |
| # dist.init_process_group("gloo", rank=int(os.environ.get("NODE_RANK")), world_size=trainer_args["num_nodes"] ) | |
| try: | |
| print(trainer_args) | |
| trainer = pl.Trainer( **trainer_args ) | |
| trainer.fit(pl_model, dataModule) | |
| res = trainer.test(pl_model, dataModule.test_dataloader(), verbose=True)[0] | |
| r, __ = InferencePlotter(pl_module=pl_model).pred_and_plot( mode="test") | |
| res["test_pearson-r"] = r | |
| print(res) | |
| finally: | |
| if args['main'].get("num_nodes", None): | |
| # dist.destroy_process_group() | |
| pass |