|
""" |
|
Transformer implementation adapted from CLIP ViT: |
|
https://github.com/openai/CLIP/blob/4c0275784d6d9da97ca1f47eaaee31de1867da91/clip/model.py |
|
""" |
|
|
|
import math |
|
|
|
import torch as th |
|
import torch.nn as nn |
|
|
|
|
|
def convert_module_to_f16(l): |
|
""" |
|
Convert primitive modules to float16. |
|
""" |
|
if isinstance(l, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)): |
|
l.weight.data = l.weight.data.half() |
|
if l.bias is not None: |
|
l.bias.data = l.bias.data.half() |
|
|
|
|
|
class LayerNorm(nn.LayerNorm): |
|
""" |
|
Implementation that supports fp16 inputs but fp32 gains/biases. |
|
""" |
|
|
|
def forward(self, x: th.Tensor): |
|
return super().forward(x.float()).to(x.dtype) |
|
|
|
|
|
class MultiheadAttention(nn.Module): |
|
def __init__(self, n_ctx, width, heads): |
|
super().__init__() |
|
self.n_ctx = n_ctx |
|
self.width = width |
|
self.heads = heads |
|
self.c_qkv = nn.Linear(width, width * 3) |
|
self.c_proj = nn.Linear(width, width) |
|
self.attention = QKVMultiheadAttention(heads, n_ctx) |
|
|
|
def forward(self, x): |
|
x = self.c_qkv(x) |
|
x = self.attention(x) |
|
x = self.c_proj(x) |
|
return x |
|
|
|
|
|
class MLP(nn.Module): |
|
def __init__(self, width): |
|
super().__init__() |
|
self.width = width |
|
self.c_fc = nn.Linear(width, width * 4) |
|
self.c_proj = nn.Linear(width * 4, width) |
|
self.gelu = nn.GELU() |
|
|
|
def forward(self, x): |
|
return self.c_proj(self.gelu(self.c_fc(x))) |
|
|
|
|
|
class QKVMultiheadAttention(nn.Module): |
|
def __init__(self, n_heads: int, n_ctx: int): |
|
super().__init__() |
|
self.n_heads = n_heads |
|
self.n_ctx = n_ctx |
|
|
|
def forward(self, qkv): |
|
bs, n_ctx, width = qkv.shape |
|
attn_ch = width // self.n_heads // 3 |
|
scale = 1 / math.sqrt(math.sqrt(attn_ch)) |
|
qkv = qkv.view(bs, n_ctx, self.n_heads, -1) |
|
q, k, v = th.split(qkv, attn_ch, dim=-1) |
|
weight = th.einsum( |
|
"bthc,bshc->bhts", q * scale, k * scale |
|
) |
|
wdtype = weight.dtype |
|
weight = th.softmax(weight.float(), dim=-1).type(wdtype) |
|
return th.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1) |
|
|
|
|
|
class ResidualAttentionBlock(nn.Module): |
|
def __init__( |
|
self, |
|
n_ctx: int, |
|
width: int, |
|
heads: int, |
|
): |
|
super().__init__() |
|
|
|
self.attn = MultiheadAttention( |
|
n_ctx, |
|
width, |
|
heads, |
|
) |
|
self.ln_1 = LayerNorm(width) |
|
self.mlp = MLP(width) |
|
self.ln_2 = LayerNorm(width) |
|
|
|
def forward(self, x: th.Tensor): |
|
x = x + self.attn(self.ln_1(x)) |
|
x = x + self.mlp(self.ln_2(x)) |
|
return x |
|
|
|
|
|
class Transformer(nn.Module): |
|
def __init__( |
|
self, |
|
n_ctx: int, |
|
width: int, |
|
layers: int, |
|
heads: int, |
|
): |
|
super().__init__() |
|
self.n_ctx = n_ctx |
|
self.width = width |
|
self.layers = layers |
|
self.resblocks = nn.ModuleList( |
|
[ |
|
ResidualAttentionBlock( |
|
n_ctx, |
|
width, |
|
heads, |
|
) |
|
for _ in range(layers) |
|
] |
|
) |
|
|
|
def forward(self, x: th.Tensor): |
|
for block in self.resblocks: |
|
x = block(x) |
|
return x |
|
|
