chendl/compositional_test

1# Copyright 2020 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15# This test is meant to be run in on an instance with TPUs like this:16#17#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py18#19# Replace 8 with the number of TPU cores you have.20#21 22import sys23from typing import Dict24 25from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available26from transformers.utils import logging27 28 29logger = logging.get_logger(__name__)30 31 32if is_torch_available():33    import torch34    from torch import nn35    from torch.utils.data import Dataset36 37    from transformers import Trainer38 39    class DummyDataset(Dataset):40        def __init__(self, length: int = 101):41            self.length = length42 43        def __len__(self):44            return self.length45 46        def __getitem__(self, i) -> int:47            return i48 49    class DummyDataCollator:50        def __call__(self, features):51            return {"input_ids": torch.tensor(features), "labels": torch.tensor(features)}52 53    class DummyModel(nn.Module):54        def __init__(self):55            super().__init__()56            # Add some (unused) params otherwise DDP will complain.57            self.fc = nn.Linear(120, 80)58 59        def forward(self, input_ids, labels=None):60            if labels is not None:61                return torch.tensor(0.0, device=input_ids.device), input_ids62            else:63                return input_ids64 65 66def main():67    parser = HfArgumentParser((TrainingArguments,))68    sys.argv += ["--output_dir", "./examples"]69    training_args = parser.parse_args_into_dataclasses()[0]70 71    logger.warning(72        f"Process rank: {training_args.local_rank}, device: {training_args.device}, "73        f"tpu_num_cores: {training_args.tpu_num_cores}",74    )75 76    # Essentially, what we want to verify in the distributed case is77    # that we get all samples back, in the right order.78    # (this is crucial for prediction for instance)79    for dataset_length in [1001, 256, 15]:80        dataset = DummyDataset(dataset_length)81 82        def compute_metrics(p: EvalPrediction) -> Dict:83            sequential = list(range(len(dataset)))84            success = p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential85            return {"success": success}86 87        trainer = Trainer(88            model=DummyModel(),89            args=training_args,90            data_collator=DummyDataCollator(),91            eval_dataset=dataset,92            compute_metrics=compute_metrics,93        )94        metrics = trainer.evaluate()95        logger.info(metrics)96        if metrics["eval_success"] is not True:97            logger.error(metrics)98            exit(1)99 100        p = trainer.predict(dataset)101        logger.info(p.metrics)102        if p.metrics["test_success"] is not True:103            logger.error(p.metrics)104            exit(1)105 106        trainer.args.eval_accumulation_steps = 2107 108        metrics = trainer.evaluate()109        logger.info(metrics)110        if metrics["eval_success"] is not True:111            logger.error(metrics)112            exit(1)113 114        p = trainer.predict(dataset)115        logger.info(p.metrics)116        if p.metrics["test_success"] is not True:117            logger.error(p.metrics)118            exit(1)119 120        trainer.args.eval_accumulation_steps = None121 122    logger.info("🔥 All distributed tests successful")123 124 125def _mp_fn(index):126    # For xla_spawn (TPUs)127    main()128 129 130if __name__ == "__main__":131    main()132