cleaning up demos (#313)

* cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos * cleaning up demos
2026-09-09 11:32:07 +08:00 · 2019-10-05 16:39:05 -04:00
parent f7d762416c
commit 07c5d22ae3
9 changed files with 194 additions and 99 deletions
@@ -0,0 +1,55 @@
+"""
+Multi-node example (GPU)
+"""
+import os
+import numpy as np
+import torch
+
+from argparse import ArgumentParser
+from pytorch_lightning import Trainer
+from examples.basic_examples.lightning_module_template import LightningTemplateModel
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+def main(hparams):
+    """
+    Main training routine specific for this project
+    :param hparams:
+    :return:
+    """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(hparams)
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer(
+        gpus=2,
+        nb_gpu_nodes=2,
+        distributed_backend='ddp'
+    )
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_parser = ArgumentParser(add_help=False)
+
+    # each LightningModule defines arguments relevant to it
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    hyperparams = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(hyperparams)