「Python で学ぶ画像認識」の画像分類プログラムを自前の Transformer に改修しました。

はじめに

独自の非自己回帰型の Transformer を提案しています。音声合成については、

で報告させていただきました。音声認識については

で報告させていただきました。機械翻訳については

で報告させていただきました。今回、提案している Transformer の非自己回帰型の使い方を、「Python で学ぶ画像認識」という本の 130 ページで解説されている Vision Transformer に応用して、学習を行い 64 % という正解率を得たので報告させていただきます。

改修点

Transfomer Decoder と Downsampler を付け加えた。

提案する非自己回帰型の Transformer の使い方では、Transformer Decoder の k, v入力に Transfomer Encoder の出力を、q 入力に Transformer Encoder の出力を Downsampling した計算量を用いる。Downsampling の割合 ds_rate は、0.5 とした。正解率は、通常計算が 63.8% であるのに対し、ds_rate = 0.5 の計算が 63.7% であった。ダウンサンプリングしたので、Transformer Decoder の計算量が減少し、通常計算では、CPU で 1 epoch 7分50秒かかっていた学習が、ds_rate = 0.5 の計算では、5分40秒で完了するようになった。0.72 倍である。

改修したプログラム

「Python で学ぶ画像認識」付属のプログラム、 4_classification/4_3_transformer/4_3_transformer.ipynb を改修した。Transformer Encoder のセルの下に TransformerDecoder のセルを加えた。このセルのプログラムは、この本の6章のものである。

class TransformerDecoderLayer(nn.Module):
    '''
    Transformerデコーダ層
    dim_hidden     : 特徴量次元
    num_heads      : アテンションヘッドの数
    dim_feedforward: FNNの中間特徴量次元
    dropout        : ドロップアウト確率
    '''
    def __init__(self, dim_hidden: int, num_heads: int,
                 dim_feedforward: int=2048, dropout: int=0.1):
        super().__init__()

        # 単語ベクトルに対するマルチヘッド自己アテンション
        self.self_attention = nn.MultiheadAttention(
            dim_hidden, num_heads, dropout, batch_first=True)
        self.norm1 = nn.LayerNorm(dim_hidden)
        self.dropout1 = nn.Dropout(dropout)

        # エンコーダとデコーダ中間出力に対するマルチヘッドアテンション
        self.cross_attention = nn.MultiheadAttention(
            dim_hidden, num_heads, dropout, batch_first=True)
        self.norm2 = nn.LayerNorm(dim_hidden)
        self.dropout2 = nn.Dropout(dropout)

        # FNN＆レイヤー正規化
        self.fnn = nn.Sequential(
            nn.Linear(dim_hidden, dim_feedforward),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, dim_hidden)
        )
        self.norm3 = nn.LayerNorm(dim_hidden)
        self.dropout3 = nn.Dropout(dropout)

    '''
    Transformerデコーダ層の順伝播
    tgt     : デコーダへの入力系列,
              [バッチサイズ, 系列長, 埋め込み次元]
    memory  : エンコーダ層の特徴量, [バッチサイズ, 1, 埋め込み次元]
    tgt_mask: 入力系列のマスク, [系列長, 系列長]
    '''
    def forward(self, tgt: torch.Tensor, memory: torch.Tensor,
                tgt_mask: torch.Tensor=None):

        # デコーダ入力に対するマスク付きマルチヘッド自己アテンション
        tgt2 = self.self_attention(
            query=tgt, key=tgt, value=tgt, attn_mask=tgt_mask)[0]
        tgt = tgt + self.dropout1(tgt2)  # 残差接続
        tgt = self.norm1(tgt)            # レイヤー正規化

        # エンコーダとデコーダ中間出力に対するマルチヘッド交差アテンション
        tgt2 = self.cross_attention(
            query=tgt, key=memory, value=memory)[0]
        tgt = tgt + self.dropout2(tgt2)  # 残差接続
        tgt = self.norm2(tgt)            # レイヤー正規化

        # FNN&レイヤー正規化
        tgt2 = self.fnn(tgt)
        tgt = tgt + self.dropout3(tgt2)  # 残差接続
        tgt = self.norm3(tgt)            # レイヤー正規化

        return tgt

VisionTransformer に self.ds_rate 変数と downsample 関数を追加し、TransformerDecoderLayer を追加した。

class VisionTransformer(nn.Module):
    '''
    Vision Transformer
    num_classes    : 分類対象の物体クラス数
    img_size       : 入力画像の大きさ(幅と高さ等しいことを想定)
    patch_size     : パッチの大きさ(幅と高さ等しいことを想定)
    dim_hidden     : 入力特徴量の次元
    num_heads      : マルチヘッドアテンションのヘッド数
    dim_feedforward: FNNにおける中間特徴量の次元
    num_layers     : Transformerエンコーダの層数
    '''
    def __init__(self, num_classes: int, img_size: int,
                 patch_size: int, dim_hidden: int, num_heads: int,
                 dim_feedforward: int, num_layers: int):
        super().__init__()

        # 画像をパッチに分解するために、
        # 画像の大きさがパッチの大きさで割り切れるか確認
        assert img_size % patch_size == 0

        self.img_size = img_size
        self.patch_size = patch_size
        self.ds_rate = 0.5

        # パッチの行数と列数はともにimg_size // patch_sizeであり、
        # パッチ数はその2乗になる
        num_patches = (img_size // patch_size) ** 2

        # パッチ特徴量はパッチを平坦化することにより生成されるため、
        # その次元はpatch_size * patch_size * 3 (RGBチャネル)
        dim_patch = 3 * patch_size ** 2

        # パッチ特徴量をTransformerエンコーダに入力する前に
        # パッチ特徴量の次元を変換する全結合層
        self.patch_embed = nn.Linear(dim_patch, dim_hidden)

        # 位置埋め込み(パッチ数 + クラス埋め込みの分を用意)
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, dim_hidden))

        # クラス埋め込み
        self.class_token = nn.Parameter(
            torch.zeros((1, 1, dim_hidden)))

        # Transformerエンコーダ層
        self.layers = nn.ModuleList([TransformerEncoderLayer(
            dim_hidden, num_heads, dim_feedforward
        ) for _ in range(num_layers)])

        # Transformerデコーダ層
        self.Dlayers = nn.ModuleList([TransformerDecoderLayer(
            dim_hidden, num_heads, dim_feedforward
        ) for _ in range(num_layers)])

        # ロジットを生成する前のレイヤー正規化と全結合
        self.norm = nn.LayerNorm(dim_hidden)
        self.linear = nn.Linear(dim_hidden, num_classes)

    '''
    順伝播関数
    x           : 入力, [バッチサイズ, 入力チャネル数, 高さ, 幅]
    return_embed: 特徴量を返すかロジットを返すかを選択する真偽値
    '''
    def forward(self, x: torch.Tensor, return_embed: bool=False):
        bs, c, h, w = x.shape

        # 入力画像の大きさがクラス生成時に指定したimg_sizeと
        # 合致しているか確認
        assert h == self.img_size and w == self.img_size

        # 高さ軸と幅軸をそれぞれパッチ数 * パッチの大きさに分解し、
        # [バッチサイズ, チャネル数, パッチの行数, パッチの大きさ,
        #                            パッチの列数, パッチの大きさ]
        # の形にする
        #print( "before x size:", x.size())
        x = x.view(bs, c, h // self.patch_size, self.patch_size,
                   w // self.patch_size, self.patch_size)

        # permute関数により
        # [バッチサイズ, パッチ行数, パッチ列数, チャネル,
        #                パッチの大きさ, パッチの大きさ]
        # の形にする
        x = x.permute(0, 2, 4, 1, 3, 5)

        # パッチを平坦化
        # permute関数適用後にはメモリ上のデータ配置の整合性の関係で
        # view関数を使えないのでreshape関数を使用
        x = x.reshape(
            bs, (h // self.patch_size) * (w // self.patch_size), -1)
        x = self.patch_embed(x)
        #print( "after patch_embed x size:", x.size() )
        # クラス埋め込みをバッチサイズ分用意
        class_token = self.class_token.expand(bs, -1, -1)
        x = torch.cat((class_token, x), dim=1)
        x += self.pos_embed

        # Transformerエンコーダ層を適用
        for layer in self.layers:
            x = layer(x)

        xa= self.downsample( x )

        for layer in self.Dlayers:
           xa = layer( xa, x )

        # クラス埋め込みをベースとしたの特徴量を抽出
        x = x[:, 0]
        x = self.norm(x)

        if return_embed:
            return x

        x = self.linear(x)

        return x

    '''
    モデルパラメータが保持されているデバイスを返す関数
    '''
    def get_device(self):
        return self.linear.weight.device

    '''
    モデルを複製して返す関数
    '''
    def copy(self):
        return copy.deepcopy(self)

    #def downsample(self, enc_out, input_lengths ):
    def downsample(self, enc_out ):

        max_label_length = int( round( enc_out.size(1) * self.ds_rate ) )

        polated_lengths = torch.round( torch.ones( enc_out.size(0) ) * enc_out.size(1) * self.ds_rate ).long()

        #outputs_lens = torch.ceil( input_lengths * self.ds_rate ).long()

        x = enc_out
        out_lens = polated_lengths

        for i in range( x.size(0) ):
            x0 = torch.unsqueeze( x[i], dim = 0 )
            x0 = x0.permute( 0,2,1)
            x_out = torch.nn.functional.interpolate(x0, size = (out_lens[i]), mode='nearest-exact')
            z = torch.zeros( x_out.size(0), x_out.size(1), max_label_length )
            if z.size(2) > x_out.size(2):
                z[:,:,:x_out.size(2)] = x_out[:,:,:]
            else:
                z[:,:,:] = x_out[:,:,:z.size(2)]
            x_out = z.permute( 0, 2, 1 )
            if i == 0:
                y = x_out
            if i > 0:
                y = torch.cat( (y, x_out), dim = 0 )

        #return y, outputs_lens
        return y

計算に使用したハイパーパラメータ

class Config:
    '''
    ハイパーパラメータとオプションの設定
    '''
    def __init__(self):
        self.val_ratio = 0.2       # 検証に使う学習セット内のデータの割合
        self.patch_size = 4        # パッチサイズ
        self.dim_hidden = 512      # 隠れ層の次元
        self.num_heads = 8         # マルチヘッドアテンションのヘッド数
        self.dim_feedforward = 512 # Transformerエンコーダ層内のFNNにおける隠れ層の特徴量次元
        #self.num_layers = 6        # Transformerエンコーダの層数
        self.num_layers = 3        # Transformerエンコーダの層数
        self.num_epochs = 30       # 学習エポック数
        self.lr = 1e-2             # 学習率
        self.moving_avg = 20       # 移動平均で計算する損失と正確度の値の数
        self.batch_size = 32       # バッチサイズ
        #self.num_workers = 2       # データローダに使うCPUプロセスの数
        self.num_workers = 0       # データローダに使うCPUプロセスの数
        #self.device = 'cuda'       # 学習に使うデバイス
        self.device = 'cpu'       # 学習に使うデバイス
        self.num_samples = 200     # t-SNEでプロットするサンプル数

これで計算しました。