Added pyproject.toml and more code formatting and checks

CONTRIBUTING.md

@@ -122,6 +122,12 @@ To add a new template,
  bash scripts/run_code_style.sh source_lint
  ```
 
+- To check if generated code complies with a style guide:
+
+ ```sh
+ pnpm test && sh ./scripts/run_tests.sh unzip && pnpm dist_lint
+ ```
+
 _NOTE: Even if you have a half-completed/working PR, sending a PR is still a valid contribution and we can help you finish the PR._
 
 **NOTE : When sending a PR, please kindly check if the changes are required to run in the CI.**

pyproject.toml

@@ -0,0 +1,23 @@
+[tool.black]
+line-length = 120
+target-version = ['py38', 'py39']
+include = '\.pyi?$'
+exclude = '''
+
+(
+ /(
+ \.eggs # exclude a few common directories in the
+ | \.git # root of the project
+ | \.hg
+ | \.mypy_cache
+ | \.tox
+ | \.venv
+ | _build
+ | buck-out
+ | build
+ | dist
+ )/
+ | foo.py # also separately exclude a file named foo.py in
+ # the root of the project
+)
+'''

scripts/run_code_style.sh

@@ -6,11 +6,13 @@ if [ $1 == "dist_lint" ]; then
  # Check that ./dist-tests/ exists and code is unzipped
  ls ./dist-tests/vision-classification-all/main.py
  # Comment dist-tests in .gitignore to make black running on ./dist-tests folder
- # TODO:
+ sed -i "s/dist-tests/# dist-tests/g" .gitignore
+
  ufmt diff .
  flake8 --select F401,F821 ./dist-tests # find unused imports and non imported objects
+
  # Restore .gitignore
- # TODO:
+ sed -i "s/\([# ]\+\)dist-tests/dist-tests/g" .gitignore
 elif [ $1 == "source_lint" ]; then
  ufmt diff .
 elif [ $1 == "fmt" ]; then

src/templates/template-common/utils.py

@@ -16,11 +16,7 @@
 
 #::: } :::#
 #::: if (it.save_training || it.save_evaluation) { :::#
-from ignite.handlers import (
- Checkpoint,
- DiskSaver,
- global_step_from_engine,
-) # usort: skip
+from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine # usort: skip
 
 #::: } else { :::#
 from ignite.handlers import Checkpoint
@@ -84,9 +80,7 @@ def log_metrics(engine: Engine, tag: str) -> None:
  tag
  a string to add at the start of output.
  """
- metrics_format = "{0} [{1}/{2}]: {3}".format(
- tag, engine.state.epoch, engine.state.iteration, engine.state.metrics
- )
+ metrics_format = "{0} [{1}/{2}]: {3}".format(tag, engine.state.epoch, engine.state.iteration, engine.state.metrics)
  engine.logger.info(metrics_format)
 
 
@@ -175,21 +169,13 @@ def setup_exp_logging(config, trainer, optimizers, evaluators):
  """Setup Experiment Tracking logger from Ignite."""
 
  #::: if (it.logger === 'clearml') { :::#
- logger = common.setup_clearml_logging(
- trainer, optimizers, evaluators, config.log_every_iters
- )
+ logger = common.setup_clearml_logging(trainer, optimizers, evaluators, config.log_every_iters)
  #::: } else if (it.logger === 'mlflow') { :::#
- logger = common.setup_mlflow_logging(
- trainer, optimizers, evaluators, config.log_every_iters
- )
+ logger = common.setup_mlflow_logging(trainer, optimizers, evaluators, config.log_every_iters)
  #::: } else if (it.logger === 'neptune') { :::#
- logger = common.setup_neptune_logging(
- trainer, optimizers, evaluators, config.log_every_iters
- )
+ logger = common.setup_neptune_logging(trainer, optimizers, evaluators, config.log_every_iters)
  #::: } else if (it.logger === 'polyaxon') { :::#
- logger = common.setup_plx_logging(
- trainer, optimizers, evaluators, config.log_every_iters
- )
+ logger = common.setup_plx_logging(trainer, optimizers, evaluators, config.log_every_iters)
  #::: } else if (it.logger === 'tensorboard') { :::#
  logger = common.setup_tb_logging(
  config.output_dir,
@@ -199,13 +185,9 @@ def setup_exp_logging(config, trainer, optimizers, evaluators):
  config.log_every_iters,
  )
  #::: } else if (it.logger === 'visdom') { :::#
- logger = common.setup_visdom_logging(
- trainer, optimizers, evaluators, config.log_every_iters
- )
+ logger = common.setup_visdom_logging(trainer, optimizers, evaluators, config.log_every_iters)
  #::: } else if (it.logger === 'wandb') { :::#
- logger = common.setup_wandb_logging(
- trainer, optimizers, evaluators, config.log_every_iters
- )
+ logger = common.setup_wandb_logging(trainer, optimizers, evaluators, config.log_every_iters)
  #::: } :::#
  return logger
 

src/templates/template-text-classification/data.py

@@ -48,21 +48,12 @@ def setup_data(config):
  if local_rank > 0:
  idist.barrier()
  #::: } :::#
-
- dataset_train, dataset_eval = load_dataset(
- "imdb", split=["train", "test"], cache_dir=config.data_path
- )
- tokenizer = AutoTokenizer.from_pretrained(
- config.model, cache_dir=config.tokenizer_dir, do_lower_case=True
- )
+ dataset_train, dataset_eval = load_dataset("imdb", split=["train", "test"], cache_dir=config.data_path)
+ tokenizer = AutoTokenizer.from_pretrained(config.model, cache_dir=config.tokenizer_dir, do_lower_case=True)
  train_texts, train_labels = dataset_train["text"], dataset_train["label"]
  test_texts, test_labels = dataset_eval["text"], dataset_eval["label"]
- dataset_train = TransformerDataset(
- train_texts, train_labels, tokenizer, config.max_length
- )
- dataset_eval = TransformerDataset(
- test_texts, test_labels, tokenizer, config.max_length
- )
+ dataset_train = TransformerDataset(train_texts, train_labels, tokenizer, config.max_length)
+ dataset_eval = TransformerDataset(test_texts, test_labels, tokenizer, config.max_length)
  #::: if (it.use_dist) { :::#
  if local_rank == 0:
  idist.barrier()

src/templates/template-text-classification/main.py

@@ -45,9 +45,7 @@ def run(local_rank: int, config: Any):
  )
 
  config.lr *= idist.get_world_size()
- optimizer = idist.auto_optim(
- optim.AdamW(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)
- )
+ optimizer = idist.auto_optim(optim.AdamW(model.parameters(), lr=config.lr, weight_decay=config.weight_decay))
  loss_fn = nn.BCEWithLogitsLoss().to(device=device)
 
  le = config.num_iters_per_epoch
@@ -56,9 +54,7 @@ def run(local_rank: int, config: Any):
  (le * config.num_warmup_epochs, config.lr),
  (le * config.max_epochs, 0.0),
  ]
- lr_scheduler = PiecewiseLinear(
- optimizer, param_name="lr", milestones_values=milestones_values
- )
+ lr_scheduler = PiecewiseLinear(optimizer, param_name="lr", milestones_values=milestones_values)
 
  # setup metrics to attach to evaluator
  metrics = {
@@ -67,9 +63,7 @@ def run(local_rank: int, config: Any):
  }
 
  # trainer and evaluator
- trainer = setup_trainer(
- config, model, optimizer, loss_fn, device, dataloader_train.sampler
- )
+ trainer = setup_trainer(config, model, optimizer, loss_fn, device, dataloader_train.sampler)
  evaluator = setup_evaluator(config, model, metrics, device)
 
  # setup engines logger with python logging
@@ -98,9 +92,7 @@ def run(local_rank: int, config: Any):
  #::: } else { :::#
  to_save_eval = None
  #::: } :::#
- ckpt_handler_train, ckpt_handler_eval = setup_handlers(
- trainer, evaluator, config, to_save_train, to_save_eval
- )
+ ckpt_handler_train, ckpt_handler_eval = setup_handlers(trainer, evaluator, config, to_save_train, to_save_eval)
  #::: } else if (it.patience || it.terminate_on_nan || it.limit_sec) { :::#
  setup_handlers(trainer, evaluator, config)
  #::: } :::#

src/templates/template-text-classification/trainers.py

@@ -22,15 +22,9 @@ def setup_trainer(
 
  def train_function(engine: Union[Engine, DeterministicEngine], batch: Any):
  input_ids = batch["input_ids"].to(device, non_blocking=True, dtype=torch.long)
- attention_mask = batch["attention_mask"].to(
- device, non_blocking=True, dtype=torch.long
- )
- token_type_ids = batch["token_type_ids"].to(
- device, non_blocking=True, dtype=torch.long
- )
- labels = (
- batch["label"].view(-1, 1).to(device, non_blocking=True, dtype=torch.float)
- )
+ attention_mask = batch["attention_mask"].to(device, non_blocking=True, dtype=torch.long)
+ token_type_ids = batch["token_type_ids"].to(device, non_blocking=True, dtype=torch.long)
+ labels = batch["label"].view(-1, 1).to(device, non_blocking=True, dtype=torch.float)
 
  model.train()
 
@@ -75,15 +69,9 @@ def evalutate_function(engine: Engine, batch: Any):
  model.eval()
 
  input_ids = batch["input_ids"].to(device, non_blocking=True, dtype=torch.long)
- attention_mask = batch["attention_mask"].to(
- device, non_blocking=True, dtype=torch.long
- )
- token_type_ids = batch["token_type_ids"].to(
- device, non_blocking=True, dtype=torch.long
- )
- labels = (
- batch["label"].view(-1, 1).to(device, non_blocking=True, dtype=torch.float)
- )
+ attention_mask = batch["attention_mask"].to(device, non_blocking=True, dtype=torch.long)
+ token_type_ids = batch["token_type_ids"].to(device, non_blocking=True, dtype=torch.long)
+ labels = batch["label"].view(-1, 1).to(device, non_blocking=True, dtype=torch.float)
 
  with autocast(enabled=config.use_amp):
  output = model(input_ids, attention_mask, token_type_ids)

src/templates/template-text-classification/utils.py

@@ -46,6 +46,7 @@ def setup_handlers(
  #::: } :::#
 
  #::: if (it.patience) { :::#
+
  # early stopping
  def score_fn(engine: Engine):
  return engine.state.metrics["Accuracy"]

src/templates/template-vision-classification/main.py

@@ -43,9 +43,7 @@ def run(local_rank: int, config: Any):
  lr_scheduler = PiecewiseLinear(optimizer, "lr", milestones_values=milestones_values)
 
  # trainer and evaluator
- trainer = setup_trainer(
- config, model, optimizer, loss_fn, device, dataloader_train.sampler
- )
+ trainer = setup_trainer(config, model, optimizer, loss_fn, device, dataloader_train.sampler)
  evaluator = setup_evaluator(config, model, device)
 
  # attach metrics to evaluator
@@ -84,9 +82,7 @@ def run(local_rank: int, config: Any):
  #::: } else { :::#
  to_save_eval = None
  #::: } :::#
- ckpt_handler_train, ckpt_handler_eval = setup_handlers(
- trainer, evaluator, config, to_save_train, to_save_eval
- )
+ ckpt_handler_train, ckpt_handler_eval = setup_handlers(trainer, evaluator, config, to_save_train, to_save_eval)
  #::: } else if (it.patience || it.terminate_on_nan || it.limit_sec) { :::#
  setup_handlers(trainer, evaluator, config)
  #::: } :::#

src/templates/template-vision-classification/test_all.py

@@ -23,9 +23,7 @@ def set_up():
 
 @pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
 def test_setup_data():
- config = Namespace(
- data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
- )
+ config = Namespace(data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0)
  dataloader_train, dataloader_eval = setup_data(config)
 
  assert isinstance(dataloader_train, DataLoader)

src/templates/template-vision-classification/utils.py

@@ -46,6 +46,7 @@ def setup_handlers(
  #::: } :::#
 
  #::: if (it.patience) { :::#
+
  # early stopping
  def score_fn(engine: Engine):
  return -engine.state.metrics["eval_loss"]

src/templates/template-vision-dcgan/main.py

@@ -49,12 +49,8 @@ def run(local_rank: int, config: Any):
  loss_fn = nn.BCELoss().to(device=device)
 
  # optimizers
- optimizer_d = idist.auto_optim(
- optim.Adam(model_d.parameters(), lr=config.lr, betas=(0.5, 0.999))
- )
- optimizer_g = idist.auto_optim(
- optim.Adam(model_g.parameters(), lr=config.lr, betas=(0.5, 0.999))
- )
+ optimizer_d = idist.auto_optim(optim.Adam(model_d.parameters(), lr=config.lr, betas=(0.5, 0.999)))
+ optimizer_g = idist.auto_optim(optim.Adam(model_g.parameters(), lr=config.lr, betas=(0.5, 0.999)))
 
  # trainer and evaluator
  trainer = setup_trainer(
@@ -100,9 +96,7 @@ def run(local_rank: int, config: Any):
  #::: } else { :::#
  to_save_train = None
  #::: } :::#
- ckpt_handler_train, ckpt_handler_eval = setup_handlers(
- trainer, evaluator, config, to_save_train, to_save_eval
- )
+ ckpt_handler_train, ckpt_handler_eval = setup_handlers(trainer, evaluator, config, to_save_train, to_save_eval)
  #::: } else if (it.patience || it.terminate_on_nan || it.limit_sec) { :::#
  setup_handlers(trainer, evaluator, config)
  #::: } :::#

src/templates/template-vision-dcgan/test_all.py

@@ -24,9 +24,7 @@ def set_up():
 
 @pytest.mark.skipif(os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests")
 def test_setup_data():
- config = Namespace(
- data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
- )
+ config = Namespace(data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0)
  dataloader_train, dataloader_eval, _ = setup_data(config)
 
  assert isinstance(dataloader_train, DataLoader)
@@ -61,8 +59,6 @@ def test_models():
 def test_setup_trainer():
  model, optimizer, device, loss_fn, batch = set_up()
  config = Namespace(use_amp=False, train_batch_size=2, z_dim=100)
- trainer = setup_trainer(
- config, model, model, optimizer, optimizer, loss_fn, device, None
- )
+ trainer = setup_trainer(config, model, model, optimizer, optimizer, loss_fn, device, None)
  trainer.run([batch, batch])
  assert isinstance(trainer.state.output, dict)

src/templates/template-vision-dcgan/trainers.py

@@ -23,9 +23,7 @@ def setup_trainer(
 
  real_labels = torch.ones(config.train_batch_size // ws, device=device)
  fake_labels = torch.zeros(config.train_batch_size // ws, device=device)
- noise = torch.randn(
- config.train_batch_size // ws, config.z_dim, 1, 1, device=device
- )
+ noise = torch.randn(config.train_batch_size // ws, config.z_dim, 1, 1, device=device)
 
  def train_function(engine: Union[Engine, DeterministicEngine], batch: Any):
  model_g.train()

src/templates/template-vision-dcgan/utils.py

@@ -46,6 +46,7 @@ def setup_handlers(
  #::: } :::#
 
  #::: if (it.patience) { :::#
+
  # early stopping
  def score_fn(engine: Engine):
  return -engine.state.metrics["errD"]

src/templates/template-vision-segmentation/data.py

@@ -88,9 +88,7 @@ def setup_data(config: Namespace):
  "Dataset not found. You can use `download_datasets` from data.py function to download it."
  ) from e
 
- dataset_eval = VOCSegmentationPIL(
- root=config.data_path, year="2012", image_set="val", download=False
- )
+ dataset_eval = VOCSegmentationPIL(root=config.data_path, year="2012", image_set="val", download=False)
 
  val_img_size = 513
  train_img_size = 480
@@ -100,9 +98,7 @@ def setup_data(config: Namespace):
 
  transform_train = A.Compose(
  [
- A.RandomScale(
- scale_limit=(0.0, 1.5), interpolation=cv2.INTER_LINEAR, p=1.0
- ),
+ A.RandomScale(scale_limit=(0.0, 1.5), interpolation=cv2.INTER_LINEAR, p=1.0),
  A.PadIfNeeded(val_img_size, val_img_size, border_mode=cv2.BORDER_CONSTANT),
  A.RandomCrop(train_img_size, train_img_size),
  A.HorizontalFlip(),
@@ -176,7 +172,6 @@ def download_datasets(data_path):
  # Ensure that only rank 0 download the dataset
  idist.barrier()
  #::: } :::#
-
  VOCSegmentation(data_path, image_set="train", download=True)
  VOCSegmentation(data_path, image_set="val", download=True)
  #::: if (it.use_dist) { :::#

src/templates/template-vision-segmentation/main.py

@@ -66,9 +66,7 @@ def run(local_rank: int, config: Any):
  metrics = {"IoU": IoU(cm_metric), "mIoU_bg": mIoU(cm_metric)}
 
  # trainer and evaluator
- trainer = setup_trainer(
- config, model, optimizer, loss_fn, device, dataloader_train.sampler
- )
+ trainer = setup_trainer(config, model, optimizer, loss_fn, device, dataloader_train.sampler)
  evaluator = setup_evaluator(config, model, metrics, device)
 
  # setup engines logger with python logging
@@ -105,9 +103,7 @@ def run(local_rank: int, config: Any):
  #::: } else { :::#
  to_save_eval = None
  #::: } :::#
- ckpt_handler_train, ckpt_handler_eval = setup_handlers(
- trainer, evaluator, config, to_save_train, to_save_eval
- )
+ ckpt_handler_train, ckpt_handler_eval = setup_handlers(trainer, evaluator, config, to_save_train, to_save_eval)
  #::: } else if (it.patience || it.terminate_on_nan || it.limit_sec) { :::#
  setup_handlers(trainer, evaluator, config)
  #::: } :::#