def test_parse_args_parsing_gpus(monkeypatch, cli_args, expected_gpu):
"""Test parsing of gpus and instantiation of Trainer."""
monkeypatch.setattr("torch.cuda.device_count", lambda: 2)
cli_args = cli_args.split(' ') if cli_args else []
parser = LightningArgumentParser(add_help=False, parse_as_dict=False)
parser.add_lightning_class_args(Trainer, None)
with mock.patch("sys.argv", ["any.py"] + cli_args):
args = parser.parse_args()
trainer = Trainer.from_argparse_args(args)
assert trainer.data_parallel_device_ids == expected_gpu
def test_parse_args_parsing_complex_types(cli_args, expected, instantiate):
"""Test parsing complex types."""
parser = LightningArgumentParser(add_help=False, parse_as_dict=False)
parser.add_lightning_class_args(Trainer, None)
with mock.patch("sys.argv", ["any.py"] + cli_args):
args = parser.parse_args()
for k, v in expected.items():
assert getattr(args, k) == v
if instantiate:
assert Trainer.from_argparse_args(args)
def test_parse_args_parsing(cli_args, expected):
"""Test parsing simple types and None optionals not modified."""
cli_args = cli_args.split(' ') if cli_args else []
parser = LightningArgumentParser(add_help=False, parse_as_dict=False)
parser.add_lightning_class_args(Trainer, None)
with mock.patch("sys.argv", ["any.py"] + cli_args):
args = parser.parse_args()
for k, v in expected.items():
assert getattr(args, k) == v
if 'tpu_cores' not in expected or _TPU_AVAILABLE:
assert Trainer.from_argparse_args(args)
def test_add_argparse_args_redefined_error(cli_args, monkeypatch):
"""Asserts error raised in case of passing not default cli arguments."""
class _UnkArgError(Exception):
pass
def _raise():
raise _UnkArgError
parser = LightningArgumentParser(add_help=False, parse_as_dict=False)
parser.add_lightning_class_args(Trainer, None)
monkeypatch.setattr(parser, "exit", lambda *args: _raise(), raising=True)
with pytest.raises(_UnkArgError):
parser.parse_args(cli_args)
def test_add_argparse_args_redefined(cli_args):
"""Redefines some default Trainer arguments via the cli and tests the Trainer initialization correctness."""
parser = LightningArgumentParser(add_help=False, parse_as_dict=False)
parser.add_lightning_class_args(Trainer, None)
args = parser.parse_args(cli_args)
# make sure we can pickle args
pickle.dumps(args)
# Check few deprecated args are not in namespace:
for depr_name in ("gradient_clip", "nb_gpu_nodes", "max_nb_epochs"):
assert depr_name not in args
trainer = Trainer.from_argparse_args(args=args)
pickle.dumps(trainer)
assert isinstance(trainer, Trainer)
import pytorch_lightning as pl
from pytorch_lightning.utilities.cli import LightningArgumentParser
from pytorch_lightning.loggers import TensorBoardLogger
from argparse import Namespace
import os
from src.data import ContentStyleDataModule
from src.model import ASTModel
# Parse arguments
parser = LightningArgumentParser()
parser.add_lightning_class_args(pl.Trainer, None)
parser.add_lightning_class_args(ASTModel, "model")
parser.add_lightning_class_args(ContentStyleDataModule, "data")
parser.add_argument("--output_path",
type=str,
help="Directory to save outputs.",
default="output/")
parser.add_argument("--experiment_name",
type=str,
help="Name of experiment.",
default="default")
args = parser.parse_args()
# Define loggers
tb_logger = TensorBoardLogger(save_dir=args["output_path"],
name=args["experiment_name"])
# Setup model
dm = ContentStyleDataModule(**args["data"])
model = ASTModel(**args["model"])