boom boom

examples/lightning_base.py

@@ -119,7 +119,7 @@ class BaseTransformer(pl.LightningModule):
     def get_lr_scheduler(self):
         get_schedule_func = arg_to_scheduler[self.hparams.lr_scheduler]
         scheduler = get_schedule_func(
-            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps
+            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps()
         )
         scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
         return scheduler
@@ -159,19 +159,20 @@ class BaseTransformer(pl.LightningModule):
     def test_epoch_end(self, outputs):
         return self.validation_end(outputs)
 
-    @property
     def total_steps(self) -> int:
         """The number of total training steps that will be run. Used for lr scheduler purposes."""
         num_devices = max(1, self.hparams.gpus)  # TODO: consider num_tpu_cores
         effective_batch_size = self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
-        dataset_size = len(self.train_loader.dataset)
+        return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
-        return (dataset_size / effective_batch_size) * self.hparams.max_epochs
 
     def setup(self, mode):
-        if mode == "fit":
+        if mode == "test":
+            self.dataset_size = len(self.test_dataloader().dataset)
+        else:
             self.train_loader = self.get_dataloader("train", self.hparams.train_batch_size, shuffle=True)
+            self.dataset_size = len(self.train_loader.dataset)
 
-    def get_dataloader(self, type_path, batch_size, shuffle=False):
+    def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False):
         raise NotImplementedError("You must implement this for your task")
 
     def train_dataloader(self):

examples/requirements.txt

@@ -5,7 +5,7 @@ psutil
 sacrebleu
 rouge-score
 tensorflow_datasets
-pytorch-lightning==0.8.5
+pytorch-lightning==0.9.0
 matplotlib
 git-python==1.0.3
 faiss

examples/seq2seq/README.md

@@ -227,7 +227,9 @@ python run_eval.py sshleifer/distilbart-cnn-12-6 $DATA_DIR/val.source dbart_val_
     --fp16 \
     --bs 32
 ```
-### Multi-GPU Evalulation
+
+
+### Multi-GPU Evaluation
 here is a command to run xsum evaluation on 8 GPUS. It is more than linearly faster than run_eval.py in some cases 
 because it uses SortishSampler to minimize padding. You can also use it on 1 GPU. `data_dir` must have 
 `{type_path}.source` and `{type_path}.target`. Run `python run_distributed_eval.py --help` for all clargs.
@@ -353,6 +355,17 @@ runtime: 13H on V-100 16GB GPU.
 pytest examples/seq2seq/
 ```
 
+### Converting pytorch-lightning checkpoints
+pytorch lightning ``-do_predict`` often fails, after you are done training, the best way to evaluate your model is to convert it.
+
+This should be done for you, with a file called `{save_dir}/best_tfmr`. 
+
+If that file doesn't exist but you have a lightning `.ckpt` file, you can run
+```bash
+python convert_pl_checkpoint_to_hf.py PATH_TO_CKPT  randomly_initialized_hf_model_path save_dir/best_tfmr
+```
+Then either `run_eval` or `run_distributed_eval` with `save_dir/best_tfmr` (see previous sections)
+
 
 ## Experimental Features 
 These features are harder to use and not always useful.
@@ -381,4 +394,3 @@ uses 12,723 batches of length 48 and takes slightly more time 9.5 minutes.
 The feature is still experimental, because:
 + we can make it much more robust if we have memory mapped/preprocessed datasets.
 + The speedup over sortish sampler is not that large at the moment.
-

examples/seq2seq/distillation.py

@@ -13,7 +13,6 @@ from torch.nn import functional as F
 from finetune import SummarizationModule, TranslationModule
 from finetune import main as ft_main
 from initialization_utils import copy_layers, init_student
-from lightning_base import generic_train
 from transformers import AutoModelForSeq2SeqLM, MBartTokenizer, T5Config, T5ForConditionalGeneration
 from transformers.modeling_bart import shift_tokens_right
 from utils import (
@@ -22,7 +21,6 @@ from utils import (
     calculate_bleu,
     freeze_params,
     label_smoothed_nll_loss,
-    pickle_load,
     use_task_specific_params,
 )
 
@@ -412,30 +410,6 @@ def create_module(args):
     return model
 
 
-def evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
-    # TODO(SS): DELETE? Better to convert_pl_ckpt_to_hf and run_eval.py
-    exp_dir = ckpt_path.parent
-    if dest_dir is None:
-        dest_dir = exp_dir
-    clash = list(dest_dir.glob("test_generations*"))
-    if clash:
-        print(f"SKIPPING to avoid overwriting {clash}")
-    ckpt = torch.load(ckpt_path, map_location="cpu")
-    if "hparams" in ckpt:
-        args = argparse.Namespace(**ckpt["hparams"])
-    else:
-        args = argparse.Namespace(**pickle_load(exp_dir / "hparams.pkl"))
-    args.resume_from_checkpoint = str(ckpt_path)
-    args.do_train = False
-    args.output_dir = str(dest_dir)
-    args.n_gpu = 1
-    args.eval_batch_size = 16
-    Path(args.output_dir).mkdir(exist_ok=True)
-    model = create_module(args)
-    trainer: pl.Trainer = generic_train(model, args, early_stopping_callback=False)
-    trainer.test(model)
-
-
 LAYERS_TO_COPY = {
     # maps  num layers in student -> which teacher layers to copy.
     # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP

examples/seq2seq/finetune.py

@@ -36,7 +36,7 @@ from utils import (
 
 
 logger = logging.getLogger(__name__)
-
+from pytorch_lightning.utilities import rank_zero_only
 
 class SummarizationModule(BaseTransformer):
     mode = "summarization"
@@ -183,6 +183,20 @@ class SummarizationModule(BaseTransformer):
         losses.update(generative_metrics)
         all_metrics = {f"{prefix}_avg_{k}": x for k, x in losses.items()}
         all_metrics["step_count"] = self.step_count
+
+        def get_date_str(seconds=True) -> str:
+            """Returns 2019-09-25-10:02:07, for example."""
+            if seconds:
+                return time.strftime('%Y-%m-%d-%H:%M:%S')
+            else:
+                return time.strftime('%Y-%m-%d-%H:%M')
+
+        all_metrics['Time'] = get_date_str(seconds=True)
+        #all_metrics['n_obs'] =
+
+
+        all_metrics['rank'] = getattr(self.train_dataloader(), '_rank', -1.)
+
         self.save_metrics(all_metrics, prefix)  # writes to self.metrics_save_path
         preds = flatten_list([x["preds"] for x in outputs])
         return {

examples/seq2seq/test_seq2seq_examples.py

@@ -13,7 +13,7 @@ import torch
 
 import lightning_base
 from convert_pl_checkpoint_to_hf import convert_pl_to_hf
-from distillation import distill_main, evaluate_checkpoint
+from distillation import distill_main
 from finetune import SummarizationModule, main
 from run_eval import generate_summaries_or_translations, run_generate
 from run_eval_search import run_search
@@ -177,7 +177,6 @@ class TestSummarizationDistiller(unittest.TestCase):
         generate_summaries_or_translations(examples, out_path, str(model.output_dir / "best_tfmr"))
         self.assertTrue(Path(out_path).exists())
 
-        evaluate_checkpoint(ckpts[0], dest_dir=Path(tempfile.mkdtemp()))
         out_path_new = tempfile.mkdtemp()
         convert_pl_to_hf(ckpts[0], transformer_ckpts[0].parent, out_path_new)
         assert os.path.exists(os.path.join(out_path_new, "pytorch_model.bin"))
@@ -226,8 +225,6 @@ class TestSummarizationDistiller(unittest.TestCase):
         assert len(all_files) > 2
         self.assertEqual(len(transformer_ckpts), 2)
 
-        evaluate_checkpoint(ckpts[0], dest_dir=Path(tempfile.mkdtemp()))
-
     @unittest.skip("T5 distillation is broken at the moment")
     def test_distill_t5(self):
         updates = dict(

examples/seq2seq/utils.py

@@ -276,10 +276,12 @@ class DistributedSortishSampler(Sampler):
             if not dist.is_available():
                 raise RuntimeError("Requires distributed package to be available")
             num_replicas = dist.get_world_size()
+            self._rank = -1
         if rank is None:
             if not dist.is_available():
                 raise RuntimeError("Requires distributed package to be available")
             rank = dist.get_rank()
+            self._rank = rank
         self.dataset = dataset
         self.num_replicas = num_replicas
         self.rank = rank