g mlp pytorch

gMLP - Pytorch

Pytorch での gMLP (Transformer の完全な MLP 代替品) の実装

インストール

$ pip install g-mlp-pytorch

使用法

マスクされた言語モデリングの場合

 import torch
from torch import nn
from g_mlp_pytorch import gMLP

model = gMLP (
    num_tokens = 20000 ,
    dim = 512 ,
    depth = 6 ,
    seq_len = 256 ,
    circulant_matrix = True ,      # use circulant weight matrix for linear increase in parameters in respect to sequence length
    act = nn . Tanh ()               # activation for spatial gate (defaults to identity)
)

x = torch . randint ( 0 , 20000 , ( 1 , 256 ))
logits = model ( x ) # (1, 256, 20000)

画像分類用

 import torch
from g_mlp_pytorch import gMLPVision

model = gMLPVision (
    image_size = 256 ,
    patch_size = 16 ,
    num_classes = 1000 ,
    dim = 512 ,
    depth = 6
)

img = torch . randn ( 1 , 3 , 256 , 256 )
logits = model ( img ) # (1, 1000)

論文ではaMLPとして言及されているように、キーワードattn_dimを 1 つ追加することで、パフォーマンスを向上させるために少量の注意 (片頭) を追加することもできます。これはgMLPVisionとgMLPの両方に当てはまります

 import torch
from g_mlp_pytorch import gMLPVision

model = gMLPVision (
    image_size = 256 ,
    patch_size = 16 ,
    num_classes = 1000 ,
    dim = 512 ,
    depth = 6 ,
    attn_dim = 64
)

img = torch . randn ( 1 , 3 , 256 , 256 )
pred = model ( img ) # (1, 1000)

非正方形の画像とパッチのサイズ

 import torch
from g_mlp_pytorch import gMLPVision

model = gMLPVision (
    image_size = ( 256 , 128 ),
    patch_size = ( 16 , 8 ),
    num_classes = 1000 ,
    dim = 512 ,
    depth = 6 ,
    attn_dim = 64
)

img = torch . randn ( 1 , 3 , 256 , 128 )
pred = model ( img ) # (1, 1000) 

実験的

独立した研究者が、Zhihu のブログ投稿で、gMLP にマルチヘッド アプローチを使用することを提案しています。これを行うには、 heads 1より大きく設定するだけです。

 import torch
from torch import nn
from g_mlp_pytorch import gMLP

model = gMLP (
    num_tokens = 20000 ,
    dim = 512 ,
    depth = 6 ,
    seq_len = 256 ,
    causal = True ,
    circulant_matrix = True ,
    heads = 4 # 4 heads
)

x = torch . randint ( 0 , 20000 , ( 1 , 256 ))
logits = model ( x ) # (1, 256, 20000) 

引用

 @misc { liu2021pay ,
    title   = { Pay Attention to MLPs } , 
    author  = { Hanxiao Liu and Zihang Dai and David R. So and Quoc V. Le } ,
    year    = { 2021 } ,
    eprint  = { 2105.08050 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}

 @software { peng_bo_2021_5196578 ,
    author       = { PENG Bo } ,
    title        = { BlinkDL/RWKV-LM: 0.01 } ,
    month        = aug,
    year         = 2021 ,
    publisher    = { Zenodo } ,
    version      = { 0.01 } ,
    doi          = { 10.5281/zenodo.5196578 } ,
    url          = { https://doi.org/10.5281/zenodo.5196578%7D
}