PyTorch Ignite 0.4.8 : AI Tutorials : Intermediate – T5 モデルによる機械翻訳 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/07/2022 (0.4.8)

* 本ページは、Pytorch Ignite AI の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Tutorials : Intermediate – Machine Translation using PyTorch Ignite

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

AI Tutorials : Intermediate – T5 モデルによる機械翻訳

このチュートリアルは、PyTorch Ignite を使用して機械翻訳モデル (or 任意の他の seq2seq モデル) を訓練する方法の簡潔なイントロダクションです。このノートブックは HuggingFace からの Model, Dataset と Tokenizer を使用していて、それらは Hub の他のモデルと容易に置き換えられます。

 

Required Dependencies

%%capture
!pip install pytorch-ignite
!pip install transformers
!pip install datasets
!pip install sentencepiece

 

For TPUs

# VERSION = !curl -s https://api.github.com/repos/pytorch/xla/releases/latest | grep -Po '"tag_name": "v\K.*?(?=")'
# VERSION = VERSION[0].rstrip('.0') # remove trailing zero
# !pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-{VERSION}-cp37-cp37m-linux_x86_64.whl

 

一般的な Configuration

config 辞書を保持します、これは訓練の間に必要なパラメータをストアするために拡張して変更できます。後でこれらのパラメータを使用するときこのコードを参照できます。

このサンプルでは t5-small を使用しています、これは 60M パラメータを持ちます。t5 モデルが動作する方法はそれらはタスク固有の prefix を持つ入力を取ります。(“Translate English to German” のような) この prefix はモデルにそれがどのタスクを遂行する必要があるかを知らせます。詳細は ここ の元の論文を参照してください。

ここではステップ毎の少ない数のイテレーションと限定されたデータセット上で訓練します、これは train_dataset_length と epoch_length config を使用して変更できます。

config = {
    "seed": 216,
    "with_amp": False,
    "num_epochs": 1,
    "batch_size": 32,
    "output_path_": "/content",
    "model_name": "t5-small",
    "tokenizer_name": "t5-small",
    "freeze_encoder": False,
    "num_workers": 4,
    "weight_decay": 0.01,
    "learning_rate": 1e-4,
    "accumulation_steps": 1,
    "epoch_length": 500,
    "print_output_every": 50,
}

dataset_configs = {
    "source_language":"English",
    "source_text_id":"en",
    "target_language":"German",
    "target_text_id":"de",
    "max_length": 80,
    "train_dataset_length": -1,
    "validation_dataset_length": 100,
    "train_test_split": 0.3,
}

 

基本的なセットアップ

インポート

import warnings
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import random_split

import ignite
import ignite.distributed as idist
from ignite.contrib.engines import common
from ignite.engine import Engine, Events
from ignite.handlers import Checkpoint, global_step_from_engine
from ignite.metrics import Bleu
from ignite.utils import manual_seed, setup_logger

from datasets import load_dataset
from transformers import T5ForConditionalGeneration, AutoTokenizer

warnings.filterwarnings("ignore")

 

データの準備

このサンプルのために Hub から new_commentary データ (English – German) を使用していきます。

from datasets import load_dataset

dataset = load_dataset("news_commentary", "de-en")
dataset = dataset.shuffle(seed=config["seed"])

Reusing dataset news_commentary (/root/.cache/huggingface/datasets/news_commentary/de-en/11.0.0/cfab724ce975dc2da51cdae45302389860badc88b74db8570d561ced6004f8b4)



  0%|          | 0/1 [00:00<?, ?it/s]


Loading cached shuffled indices for dataset at /root/.cache/huggingface/datasets/news_commentary/de-en/11.0.0/cfab724ce975dc2da51cdae45302389860badc88b74db8570d561ced6004f8b4/cache-199f0b60779b6122.arrow

dataset = dataset["train"]
dataset = dataset.train_test_split(test_size=dataset_configs["train_test_split"])
train_dataset, validation_dataset = dataset["train"], dataset["test"]

print("Lengths")
print("\t Train Set - {}".format(len(train_dataset)))
print("\t Val Set - {}".format(len(validation_dataset)))

Loading cached split indices for dataset at /root/.cache/huggingface/datasets/news_commentary/de-en/11.0.0/cfab724ce975dc2da51cdae45302389860badc88b74db8570d561ced6004f8b4/cache-23d286abe396b3d4.arrow and /root/.cache/huggingface/datasets/news_commentary/de-en/11.0.0/cfab724ce975dc2da51cdae45302389860badc88b74db8570d561ced6004f8b4/cache-387687cf22f2e607.arrow


Lengths
     Train Set - 156207
     Val Set - 66946

Having a look at a dataset sample.

print("Example of a Datapoint \n")
print(train_dataset[0], "\n")

Example of a Datapoint 

{'id': '123598', 'translation': {'de': 'Nachrichtenberichte und „Analysen“ der staatlich kontrollierten Sender in Russland und Georgien, die ein negatives Image „des Feindes“ zeichnen, dienen lediglich dazu, die Kluft zwischen den ethnischen Gruppen noch zu vertiefen.', 'en': 'News reports and “analysis” by state-controlled channels in both Russia and Georgia that promote negative images of “the enemy” serve only to widen the gap between ethnic groups.'}}

 

Tokenizer

tokenizer は文字列からトークン id に入力を変換するために定義される必要があります。機械翻訳 tokenizer はソース言語とターゲット言語についての追加パラメータが必要です、詳細は ここ を参照してください。

tokenizer = AutoTokenizer.from_pretrained(config["tokenizer_name"])

 

Dataset クラス

データをトークン化して入力とターゲットを持つ辞書を返します。

データのサブセットで訓練したい場合 – dataset config の train_dataset_length と validation_dataset_length を変更します。全体の長さを取得するためにはそれらを -1 として保持してください。

class TransformerDataset(torch.utils.data.Dataset):
    def __init__(
        self, data, src_text_id, tgt_text_id, tokenizer, max_length, length_dataset
    ):
        self.data = data
        self.src_text_id = src_text_id
        self.tgt_text_id = tgt_text_id
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.length_dataset = length_dataset if length_dataset != -1 else len(self.data)

    def __getitem__(self, idx):
        # t5 models require a prefix describing the task
        task_prefix = "translate {} to {}: ".format(dataset_configs["source_language"], dataset_configs["target_language"])
        src_text = [task_prefix + str(self.data[idx]["translation"][self.src_text_id])]

        tgt_text = [str(self.data[idx]["translation"][self.tgt_text_id])]
        input_txt_tokenized = self.tokenizer(
            src_text, max_length=self.max_length, padding="max_length", truncation=True
        )

        with self.tokenizer.as_target_tokenizer():
            tgt_text_tokenized = self.tokenizer(
                tgt_text,
                max_length=self.max_length,
                padding="max_length",
                truncation=True,
            )

        # The pad token in target is replaced with -100 so that it doesn't get added to loss.
        tgt_text_tokenized = [
            [(l if l != self.tokenizer.pad_token_id else -100) for l in label]
            for label in tgt_text_tokenized.input_ids
        ]

        input_txt_tokenized.update({"tgt": tgt_text_tokenized[0]})

        batch = {
            k: torch.tensor(v).squeeze(0) for (k, v) in input_txt_tokenized.items()
        }
        return batch

    def __len__(self):
        return self.length_dataset

train_data = TransformerDataset(
    train_dataset,
    dataset_configs["source_text_id"],
    dataset_configs["target_text_id"],
    tokenizer,
    dataset_configs["max_length"],
    dataset_configs["train_dataset_length"],
)
val_data = TransformerDataset(
    validation_dataset,
    dataset_configs["source_text_id"],
    dataset_configs["target_text_id"],
    tokenizer,
    dataset_configs["max_length"],
    dataset_configs["validation_dataset_length"],
)

 

トレーナー

トレーナーは入力のバッチを取り、それらを (この場合はターゲットと一緒に) モデルに渡してそして損失を取得します。

 
混合精度

forward パスは、混合精度訓練のためには autocast コンテキストマネージャ内にラップされます。これはこのサンプルでは無効にされます、何故ならば batch_size 1 か 2 ではどのようなメモリの利点もないからです。それを有効にするには config で with_amp フラグを変更します。

勾配累積

勾配累積が実装されます、というのはそうでないと 1 のバッチサイズは noisy アップデートに繋がるからです。勾配を累積するためのステップ数を定義するには config の accumulation_steps 変数を確認してください。

トレーナーハンドラ

ハンドラが定義できてトレーナーエンジンに直接装着できます。ここではまた特別な関数を利用します : setup_common_training_handlers、これは既に定義された一般に使用される、有用なハンドラ (save_every_iters, clear_cuda_cache 等のような) を含みます。この関数について詳細を知るには、docs を ここ で参照してください。

# Create Trainer
def create_trainer(model, optimizer, with_amp, train_sampler, logger):
    device = idist.device()
    scaler = GradScaler(enabled=with_amp)

    def train_step(engine, batch):
        model.train()

        if batch["tgt"].device != device:
            batch = {
                k: v.to(device, non_blocking=True, dtype=torch.long)
                for (k, v) in batch.items()
            }

        src_ids = batch["input_ids"]
        src_attention_mask = batch["attention_mask"]
        tgt = batch["tgt"]

        with autocast(enabled=with_amp):
            y = model(input_ids=src_ids, attention_mask=src_attention_mask, labels=tgt)
            loss = y["loss"]
            loss /= config["accumulation_steps"]

        scaler.scale(loss).backward()

        if engine.state.iteration % config["accumulation_steps"] == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

        return {"batch loss": loss.item()}

    trainer = Engine(train_step)
    trainer.logger = logger

    metric_names = ["batch loss"]

    common.setup_common_training_handlers(
        trainer=trainer,
        train_sampler=train_sampler,
        output_names=metric_names,
        clear_cuda_cache=False,
        with_pbars=True,
    )
    return trainer

 

Evaluator

トレーナーと同様に検証ステップのために evaluator を作成します。ここでは (Bleu Score のような) メトリクスを計算します。このためには Bleu スコアはロジットではなくセンテンスを必要とします。それを行なうため ids_to_clean_text 関数が使用されます。

出力センテンスをプリントする頻度を変更したい場合には print_output_every フラグを変更できます。

# Let's now setup evaluator engine to perform model's validation and compute metrics
def create_evaluator(model, tokenizer, metrics, logger, tag="val"):
    device = idist.device()

    def ids_to_clean_text(generated_ids):
        gen_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
        return list(map(str.strip, gen_text))

    @torch.no_grad()
    def evaluate_step(engine, batch):
        model.eval()

        if batch["tgt"].device != device:
            batch = {
                k: v.to(device, non_blocking=True, dtype=torch.long)
                for (k, v) in batch.items()
            }

        src_ids = batch["input_ids"]
        src_attention_mask = batch["attention_mask"]
        tgt = batch["tgt"]
        if idist.get_world_size() > 1:
            y_pred = model.module.generate(input_ids=src_ids, attention_mask=src_attention_mask)
        else:   
            y_pred = model.generate(input_ids=src_ids, attention_mask=src_attention_mask)

        tgt = torch.where(tgt != -100, tgt, tokenizer.pad_token_id)

        preds = ids_to_clean_text(y_pred)
        tgt = ids_to_clean_text(tgt)
        preds = [_preds.split() for _preds in preds]
        tgt = [[_tgt.split()] for _tgt in tgt]
        
        if engine.state.iteration % config["print_output_every"] == 0:
            logger.info(f'\n Preds : {" ".join(preds[0])} \n')
            logger.info(f'\n Target : {" ".join(tgt[0][0])} \n')
        return preds, tgt

    evaluator = Engine(evaluate_step)

    for name, metric in metrics.items():
        metric.attach(evaluator, name)

    return evaluator

 

関数の初期化

ここでモデルと optimizer を初期化します。
get_dataloader は訓練と検証のために dataloader を返します。

def freeze_params(model):
    for par in model.parameters():
        par.requires_grad = False


def initialize():
    model = T5ForConditionalGeneration.from_pretrained(config["model_name"])
    lr = config["learning_rate"] * idist.get_world_size()
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if not any(nd in n for nd in no_decay)
            ],
            "weight_decay": config["weight_decay"],
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.0,
        },
    ]
    if config["freeze_encoder"]:
        freeze_params(model.get_encoder())

    model = idist.auto_model(model)
    optimizer = optim.AdamW(optimizer_grouped_parameters, lr=lr)
    optimizer = idist.auto_optim(optimizer)

    return model, optimizer

def get_dataloaders(train_dataset, val_dataset):
    # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
    train_loader = idist.auto_dataloader(
        train_dataset,
        batch_size=config["batch_size"],
        num_workers=config["num_workers"],
        shuffle=True,
        drop_last=True,
    )

    val_loader = idist.auto_dataloader(
        val_dataset,
        batch_size=2 * config["batch_size"],
        num_workers=config["num_workers"],
        shuffle=False,
    )
    return train_loader, val_loader

 

ロギングハンドラ

このステップはオプションですが、log_basic_info() に setup_logger() オブジェクトを渡して、異なるバージョン、現在の configuration、(ローカル rank により識別される) 現在のプロセスにより使用されるデバイスとバックエンド、そしてプロセス数 (world サイズ) のような総ての基本的な情報をログ記録することができます。idist (ignite.distributed) はこれを可能にするために get_local_rank(), backend(), get_world_size() 等のような幾つかのユティリティ関数を提供しています。

log_metrics_eval は評価を実行するためにメトリクスと評価時間をログ記録するために使用されます。

get_save_handler はそれが呼び出されるたびにモデルを output パスにセーブするために使用されます。

def log_metrics_eval(logger, epoch, elapsed, tag, metrics):
    metrics_output = "\n".join([f"\t{k}: {v}" for k, v in metrics.items()])
    logger.info(
        f"\nEpoch {epoch} - Evaluation time (seconds): {elapsed:.2f} - {tag} metrics:\n {metrics_output}"
    )


def log_basic_info(logger, config):
    logger.info(f"Train on CIFAR10")
    logger.info(f"- PyTorch version: {torch.__version__}")
    logger.info(f"- Ignite version: {ignite.__version__}")
    if torch.cuda.is_available():
        # explicitly import cudnn as torch.backends.cudnn can not be pickled with hvd spawning procs
        from torch.backends import cudnn

        logger.info(
            f"- GPU Device: {torch.cuda.get_device_name(idist.get_local_rank())}"
        )
        logger.info(f"- CUDA version: {torch.version.cuda}")
        logger.info(f"- CUDNN version: {cudnn.version()}")

    logger.info("\n")
    logger.info("Configuration:")
    for key, value in config.items():
        logger.info(f"\t{key}: {value}")
    logger.info("\n")

    if idist.get_world_size() > 1:
        logger.info("\nDistributed setting:")
        logger.info(f"\tbackend: {idist.backend()}")
        logger.info(f"\tworld size: {idist.get_world_size()}")
        logger.info("\n")


def get_save_handler(config):
    return config["output_path_"]

 

訓練の開始

これは主要ロジックがあるところです、i.e. ここの中から上記の総ての関数を呼び出します :

1. 基本的なセットアップ
   1. manual_seed() and setup_logger() を設定してから、総ての基本的な情報をログ記録します。
   2. データローダ、モデルと optimizer を初期化します。

2. トレーナーを作成するために上記のオブジェクトを使用します。
   

3. Evaluator
   1. Bleu() のような幾つかの関連 Ignite メトリクスを定義する。
   2. val_dataloader 上でメトリクスを計算するために evaluator: evaluator を作成します。
   3. 両者のデータローダでメトリクスを計算するために run_validation() を定義してそれらをログ記録します。そしてこの関数をエポック後に実行するため trainer に装着します。

4. 学習率とともに検証メトリクスがログ記録されるように、evaluator のためにマスタープロセスで setup_tb_logging() を使用して TensorBoard ロギングをセットアップします。
   

5. (メトリクスで Bleu() として定義された) 検証精度により 2 つの最良モデル (n_saved) をストアするために Checkpoint() オブジェクトを定義して、それを evaluator が動作するたびに実行できるように val_evaluator に装着します。
   

6. train_loader で num_epochs 訓練を試行する。
   

7. 訓練が完了したら Tensorboard logger を閉じる。

def training(local_rank):
    rank = idist.get_rank()
    manual_seed(config["seed"] + rank)
    device = idist.device()

    logger = setup_logger(name="NMT", distributed_rank=local_rank)
    log_basic_info(logger, config)

    train_loader, val_loader = get_dataloaders(train_data, val_data)
    model, optimizer = initialize()

    trainer = create_trainer(
        model, optimizer, config["with_amp"], train_loader.sampler, logger
    )

    metrics = {
        "bleu": Bleu(ngram=4, smooth="smooth1", average="micro"),
        "bleu_smooth_2": Bleu(ngram=4, smooth="smooth2", average="micro"),
    }

    evaluator = create_evaluator(
        model, tokenizer, metrics, logger, tag="val"
    )

    @trainer.on(Events.EPOCH_COMPLETED(every=1) | Events.COMPLETED | Events.STARTED)
    def run_validation(engine):
        epoch = trainer.state.epoch
        state = evaluator.run(val_loader)
        log_metrics_eval(
            logger, epoch, state.times["COMPLETED"], "Validation", state.metrics
        )

    if rank == 0:
        now = datetime.now().strftime("%Y%m%d-%H%M%S")
        folder_name = f"Translation_Model_backend-{idist.backend()}-{idist.get_world_size()}_{now}"
        output_path = Path(config["output_path_"]) / folder_name
        if not output_path.exists():
            output_path.mkdir(parents=True)

        logger.info(f"Output path: {output_path}")

        evaluators = {"val": evaluator}
        tb_logger = common.setup_tb_logging(
            config["output_path_"], trainer, optimizer, evaluators=evaluators
        )

    best_model_handler = Checkpoint(
        {"model": model},
        get_save_handler(config),
        filename_prefix="best",
        n_saved=2,
        global_step_transform=global_step_from_engine(trainer),
        score_name="val_bleu",
        score_function=Checkpoint.get_default_score_fn("bleu"),
    )
    evaluator.add_event_handler(Events.COMPLETED, best_model_handler)

    try:
        state = trainer.run(
            train_loader,
            max_epochs=config["num_epochs"],
            epoch_length=config["epoch_length"],
        )
    except Exception as e:
        logger.exception("")
        raise e

    if rank == 0:
        tb_logger.close()

 

実行

TPU で実行するためには backend を “xla-tpu” にそして nproc_per_node を 1 or 8 に変更します。

def run():
    with idist.Parallel(backend=None, nproc_per_node=None) as parallel:
        parallel.run(training)

if __name__ == '__main__':
  run()

2021-10-21 13:46:21,877 ignite.distributed.launcher.Parallel INFO: - Run '<function training at 0x7f1fbee15710>' in 1 processes
2021-10-21 13:46:21,918 NMT INFO: Train on CIFAR10
2021-10-21 13:46:21,920 NMT INFO: - PyTorch version: 1.9.0+cu111
2021-10-21 13:46:21,922 NMT INFO: - Ignite version: 0.5.0
2021-10-21 13:46:21,925 NMT INFO: - GPU Device: Tesla K80
2021-10-21 13:46:21,926 NMT INFO: - CUDA version: 11.1
2021-10-21 13:46:21,931 NMT INFO: - CUDNN version: 8005
2021-10-21 13:46:21,933 NMT INFO: 

2021-10-21 13:46:21,936 NMT INFO: Configuration:
2021-10-21 13:46:21,938 NMT INFO: 	seed: 216
2021-10-21 13:46:21,940 NMT INFO: 	with_amp: False
2021-10-21 13:46:21,943 NMT INFO: 	num_epochs: 1
2021-10-21 13:46:21,946 NMT INFO: 	batch_size: 32
2021-10-21 13:46:21,949 NMT INFO: 	output_path_: /content
2021-10-21 13:46:21,952 NMT INFO: 	model_name: t5-small
2021-10-21 13:46:21,956 NMT INFO: 	tokenizer_name: t5-small
2021-10-21 13:46:21,959 NMT INFO: 	freeze_encoder: False
2021-10-21 13:46:21,961 NMT INFO: 	num_workers: 4
2021-10-21 13:46:21,964 NMT INFO: 	weight_decay: 0.01
2021-10-21 13:46:21,968 NMT INFO: 	learning_rate: 0.0001
2021-10-21 13:46:21,972 NMT INFO: 	accumulation_steps: 1
2021-10-21 13:46:21,974 NMT INFO: 	epoch_length: 500
2021-10-21 13:46:21,976 NMT INFO: 	print_output_every: 50
2021-10-21 13:46:21,980 NMT INFO: 

2021-10-21 13:46:21,983 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<__main__.Transforme': 
	{'batch_size': 32, 'num_workers': 4, 'shuffle': True, 'drop_last': True, 'pin_memory': True}
2021-10-21 13:46:21,986 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<__main__.Transforme': 
	{'batch_size': 64, 'num_workers': 4, 'shuffle': False, 'pin_memory': True}
2021-10-21 13:46:26,245 NMT INFO: Output path: /content/Translation_Model_backend-None-1_20211021-134626
2021-10-21 13:46:26,327 NMT INFO: Engine run starting with max_epochs=1.
2021-10-21 13:46:28,533 NMT INFO: 
Epoch 0 - Evaluation time (seconds): 2.10 - Validation metrics:
    bleu: 0.10135051023993102
	bleu_smooth_2: 0.10169442246586281



100%|##########| 1/1 [00:00<?, ?it/s]



[1/500]   0%|           [00:00<?]


2021-10-21 13:52:00,975 NMT INFO: 
Epoch 1 - Evaluation time (seconds): 2.03 - Validation metrics:
    bleu: 0.10242125441879026
	bleu_smooth_2: 0.10276058920188186
2021-10-21 13:52:00,978 NMT INFO: Epoch[1] Complete. Time taken: 00:05:32
2021-10-21 13:52:03,141 NMT INFO: 
Epoch 1 - Evaluation time (seconds): 2.04 - Validation metrics:
    bleu: 0.10242125441879026
	bleu_smooth_2: 0.10276058920188186
2021-10-21 13:52:03,143 NMT INFO: Engine run complete. Time taken: 00:05:37
2021-10-21 13:52:03,267 ignite.distributed.launcher.Parallel INFO: End of run

以上