mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-21 13:20:08 +08:00
43 lines
937 B
Python
43 lines
937 B
Python
"""
|
|
Runs a model on a single node on CPU only..
|
|
"""
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
|
|
from test_tube import HyperOptArgumentParser, Experiment
|
|
from pytorch_lightning import Trainer
|
|
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
|
|
|
|
from examples.new_project_templates.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()
|
|
|
|
# ------------------------
|
|
# 3 START TRAINING
|
|
# ------------------------
|
|
trainer.fit(model)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main(hyperparams)
|