任意のビデオエンコーダとHuggingFaceのデコーダモデルをつなげて自作のビデオキャプショニングモデルを作る


[PyTorch Hub](https://pytorch.org/hub/)など[HuggingFace](https://huggingface.co/)以外から取得したビデオエンコーダとHuggingFaceのデコーダを使って，ビデオそのものを入力とするキャプショニングモデルの実装方法を紹介します．


実装方法のみを知りたい方は[実装の節](#実装)を参照してください．

# やりたいこと
- ビデオを入力してテキストを生成するビデオキャプショニングモデルを実装したい
  - video encoderでビデオを埋め込み
  - 埋め込みからdecoderでテキスト生成
- video encoderには，HuggingFaceには無いモデルを使用したい
- decoderにはHuggingFaceのモデルを使用したい
  - `generate()`でキャプション生成したい
  


# そもそも

ビデオを入力とするキャプショニングモデルを実装には，HuggingFaceの[`VisionEncoderDecoder`](https://huggingface.co/docs/transformers/model_doc/vision-encoder-decoder)モデルを使うと簡単です．以下のように，`encoder`と`decoder`に使用するモデルを（文字列で）指定するだけでインスタンスが生成できます．

```py
from transformers import VisionEncoderDecoderModel


model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(
            encoder='MCG-NJU/videomae-base-finetuned-kinetics',
            decoder='gpt2'
        )
```

`VisionEncoderDecoderModel`の詳細はこちら．


- <https://huggingface.co/docs/transformers/main/en/model_doc/vision-encoder-decoder#transformers.VisionEncoderDecoderModel>
- <https://huggingface.co/docs/transformers/model_doc/vision-encoder-decoder>


しかし，`encoder`と`decoder`は，どちらもHuggingFaceに存在するモデルしか用いることができません．そのため，[`torchhub`](https://pytorch.org/hub/)などから取ってきた他のvideo encoder（たとえば[3D ResNet](https://pytorch.org/hub/facebookresearch_pytorchvideo_resnet/)）を使いたい場合は，このクラスを使うことはできません．



### `generate`関数


HuggingFaceのencoder-decoderやdecoderモデルはテキスト生成のための`generate`関数を持っています．


この`generate`関数は「モデル自体のforward（順伝播） + テキスト生成」を行っており（実際にはテキスト生成ではなく，単語トークンid列の生成），その実体は[`GenerationMixin`](https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin)クラスに実装されていて，各decoderの親クラスがこれを継承しています．


`VisionEncoderDecoderModel`も`generate`関数を持っているため，これでテキスト生成ができます．しかし今回のように，encoderに任意のモデルを使う場合には，そもそも`VisionEncoderDecoderModel`の`encoder`に指定できないため，使うことができません．


以下では，decoderの関数の一つをオーバーロードすることで`generate`に埋め込みを渡せるようにする方法を紹介します．


`generate`関数のドキュメントはこちら（テキスト生成における探索方法の設定など，色々変えることができます）．


- <https://huggingface.co/docs/transformers/main_classes/text_generation>


## ビデオキャプショニングモデルの概要


### train


1. ビデオをvideo encoderに入力し，video埋め込みを得る
1. decoderへの入力には，このvideo埋め込みと，右に1単語シフトした正解テキストを与える（シフト後に先頭に`<BOS>`トークンが入る）
1. 次単語予測（シフトしていない正解テキストを予測）を行い，正解テキストと交差エントロピー損失を求める


1がencoderの`forward`関数で，2と3がdecoderの`foward`関数で行われます．


- 参考：[Image Captioning in Deep Learning -- What is Image Captioning?](https://towardsdatascience.com/image-captioning-in-deep-learning-9cd23fb4d8d2)


### val


1. ビデオをvideo encoderに入力し，video埋め込みを得る
1. decoderにvideo埋め込みと，BOSトークンを入力し，テキストを生成する
1. 正解テキストと比較する


1がencoderの`forward`関数で，2がdecoderの`generate`関数で行われます
(`generate`関数内部で，decoderの`forward`関数が呼び出されています)．


## 実装


### 設定


- ビデオエンコーダ : [PyTorchHub](https://pytorch.org/hub/)の[X3D_m](https://pytorch.org/hub/facebookresearch_pytorchvideo_x3d/)
- テキストデコーダ : [HuggingFace](https://huggingface.co/)の[GPT2LMHeadModel](https://huggingface.co/docs/transformers/main/en/model_doc/gpt2#transformers.GPT2LMHeadModel)


### 1. Encoderの実装


デコーダであるGPT2が受け付ける入力の次元は768なので，エンコーダの最後の全結合層で次元を768に変換します．


```python
class MyX3D(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = torch.hub.load(
            'facebookresearch/pytorchvideo', "x3d_m",
            pretrained=True,
            head_activation=None,
        )
        in_features = self.model.blocks[5].proj.in_features
        self.model.blocks[5].proj = nn.Linear(
            in_features, 768)


    def forward(self, pixel_values):
        output = self.model(pixel_values)
        return output
```


X3Dの詳細はこちら（`head_activation`などは`create_x3d`の引数に渡されます）．


- <https://github.com/facebookresearch/pytorchvideo/blob/1730313bb5eda307657f971365de7a1359e1a68e/pytorchvideo/models/x3d.py#L537>


### 2. decoderの実装（GPT2LMHeadModelを継承）


`GPT2LMHeadModel`を継承して自作のデコーダを作成します．


継承する目的は，[`prepare_inputs_for_generation()`](https://github.com/huggingface/transformers/blob/9728f1134bf1bcb822dcb46a205e19f12ff21323/src/transformers/models/gpt2/modeling_gpt2.py#L1010)をオーバーライドして変更するためです．


エンコーダの出力が，デコーダの`forward()`の[`encoder_hidden_states`](https://github.com/huggingface/transformers/blob/9728f1134bf1bcb822dcb46a205e19f12ff21323/src/transformers/models/gpt2/modeling_gpt2.py#L297)に引数として渡される（ことが仮定されている）ので，これが`generate()`に入力されても内部で`forward()`にも渡されるようになればOKのはずです（ビデオの特徴を反映したテキストが生成される）．


この`prepare_inputs_for_generation()`は`generate()`内部で呼び出される関数であり，`forward()`に渡す引数を選択して用意する役割を持っています．しかし`GPT2LMHeadModel`の実装はそうはなって**いない**ため，`encoder_hidden_states`は`forward()`に渡されず，このままではencoderの出力は利用されることなく捨てられてしまします．


そのため，次のようにして自作クラスを継承して，オーバーライドすることで解決します．


```py
from transformers import (
    GPT2LMHeadModel
)


class MyGPT2LMHeadModel(GPT2LMHeadModel):
    def __init__(self, config):
        super().__init__(config)


    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs
    ):
        model_inputs = super().prepare_inputs_for_generation(
            input_ids, past_key_values, inputs_embeds, **kwargs
        )
        encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
        model_inputs.update(
            {
                "encoder_hidden_states": encoder_hidden_states,
            }
        )
        return model_inputs
```


### 3. VideoCaptionModelの実装


最後にキャプションモデル全体の実装します．


```py
from transformers import (
    GPT2Tokenizer,
    GPT2Config,
)
import torch.nn as nn
import torch



class VideoCaptionModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.video_encoder = MyX3D()
        config = GPT2Config(
            is_encoder_decoder=False,
            add_cross_attention=True,
        )
        self.decoder = MyGPT2LMHeadModel.from_pretrained("gpt2", config=config)
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.decoder.config.pad_token_id = self.tokenizer.pad_token_id
        self.decoder.config.bos_token_id = self.tokenizer.bos_token_id
        self.decoder.config.eos_token_id = self.tokenizer.eos_token_id


    def forward(self, pixel_values, labels, decoder_attention_mask=None):
        encoder_hidden_states = torch.unsqueeze(
            self.video_encoder(pixel_values=pixel_values), 1
        )  # bs * 1 * 768
        decoder_input_ids = shift_tokens(
            labels, self.decoder.config.pad_token_id, self.decoder.config.bos_token_id
        )
        decoder_output = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_hidden_states,
        )


        logits = decoder_output.logits
        loss_fct = nn.CrossEntropyLoss()
        loss = loss_fct(
            logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1)
        )
        decoder_output.loss = loss
        return decoder_output


    def generate(self, pixel_values, max_length=20):
        encoder_hidden_states = torch.unsqueeze(
            self.video_encoder(pixel_values=pixel_values), 1
        )
        input_ids = torch.LongTensor(
            [[self.decoder.config.bos_token_id] for _ in range(pixel_values.size()[0])]
        ).to(pixel_values.device)
        generated_ids = self.decoder.generate(
            input_ids,
            encoder_hidden_states=encoder_hidden_states,
            max_new_tokens=max_length,
        )
        return generated_ids
```


トークン化されたテキストを右シフトする関数は，[`VisionEncoderDecoder`](https://huggingface.co/docs/transformers/model_doc/vision-encoder-decoder)の[`shift_tokens_right`](https://github.com/huggingface/transformers/blob/9728f1134bf1bcb822dcb46a205e19f12ff21323/src/transformers/models/vision_encoder_decoder/modeling_vision_encoder_decoder.py#L37)を参考に作成します．


```py
def shift_tokens(token_ids, pad_token_id, bos_token_id):
    shifted_ids = input_ids.new_zeros(token_ids.shape)
    shifted_ids[:, 1:] = input_ids[:, :-1].clone()
    shifted_ids[:, 0] = bos_token_id
    return shifted_ids
```


### 4. trainとval


以下は学習と評価の1バッチ分のコードの例です．`video`は[pytorchvideo](https://github.com/facebookresearch/pytorchvideo)などの既存のビデオローダが使えます．評価指標は[`rouge`](https://huggingface.co/docs/transformers/main/en/tasks/image_captioning#evaluate)などで計算します．


```py
model = VideoCaptionModel()
rouge = evaluate.load("rouge")



def train(batch, model):


    video = batch[0]  # (B, T, C, H, W): video clips with T frames
    text = batch[1]  # (B, L): texts of length L


    with torch.no_grad():
        token = model.tokenizer(text, return_tensors="pt", padding=True)


    output = model(
        pixel_values=video,
        labels=token.input_ids.to(video.device),
        decoder_attention_mask=token.attention_mask.to(video.device),
    )
    loss = output.loss
    return loss



def val(batch, model):
    video = batch[0]
    text = batch[1]


    generated_ids = model.generate(video, max_length=20)
    generated_text = model.tokenizer.batch_decode(
        generated_ids, skip_special_tokens=True
    )
    metrics = rouge.compute(predictions=generated_text, references=text)
    return metrics
```