min-dalle-test/min_dalle/models/dalle_bart_decoder.py

219 lines
7.8 KiB
Python
Raw Normal View History

2022-06-28 16:16:44 +00:00
from typing import List, Tuple
2022-06-27 15:57:56 +00:00
import torch
from torch import LongTensor, nn, FloatTensor, BoolTensor
2022-06-29 01:28:36 +00:00
torch.set_grad_enabled(False)
2022-06-27 15:57:56 +00:00
from .dalle_bart_encoder import GLU, AttentionBase
2022-06-27 15:57:56 +00:00
class DecoderCrossAttention(AttentionBase):
2022-06-27 15:57:56 +00:00
def forward(
self,
decoder_state: FloatTensor,
encoder_state: FloatTensor,
attention_mask: BoolTensor
) -> FloatTensor:
keys = self.k_proj.forward(encoder_state)
values = self.v_proj.forward(encoder_state)
queries = self.q_proj.forward(decoder_state)
return super().forward(keys, values, queries, attention_mask)
class DecoderSelfAttention(AttentionBase):
def forward(
self,
2022-06-27 15:57:56 +00:00
decoder_state: FloatTensor,
attention_state: FloatTensor,
2022-06-27 15:57:56 +00:00
attention_mask: BoolTensor,
2022-06-29 01:28:36 +00:00
token_mask: BoolTensor
2022-06-27 17:19:03 +00:00
) -> Tuple[FloatTensor, FloatTensor]:
2022-06-29 17:48:12 +00:00
keys = self.k_proj.forward(decoder_state)
values = self.v_proj.forward(decoder_state)
queries = self.q_proj.forward(decoder_state)
attention_state = torch.where(
2022-06-29 17:48:12 +00:00
token_mask[None, :, None],
2022-06-27 17:19:03 +00:00
torch.cat([keys, values]),
attention_state
2022-06-27 15:57:56 +00:00
)
batch_count = decoder_state.shape[0]
2022-06-30 11:41:31 +00:00
keys = attention_state[:batch_count]
values = attention_state[batch_count:]
2022-06-27 15:57:56 +00:00
decoder_state = super().forward(keys, values, queries, attention_mask)
return decoder_state, attention_state
2022-06-27 15:57:56 +00:00
class DecoderLayer(nn.Module):
def __init__(
self,
2022-06-27 15:57:56 +00:00
image_token_count: int,
head_count: int,
embed_count: int,
glu_embed_count: int
):
super().__init__()
self.image_token_count = image_token_count
self.pre_self_attn_layer_norm = nn.LayerNorm(embed_count)
self.self_attn = DecoderSelfAttention(head_count, embed_count)
2022-06-27 15:57:56 +00:00
self.self_attn_layer_norm = nn.LayerNorm(embed_count)
self.pre_encoder_attn_layer_norm = nn.LayerNorm(embed_count)
self.encoder_attn = DecoderCrossAttention(head_count, embed_count)
2022-06-27 15:57:56 +00:00
self.encoder_attn_layer_norm = nn.LayerNorm(embed_count)
self.glu = GLU(embed_count, glu_embed_count)
2022-06-27 15:57:56 +00:00
2022-06-29 01:28:36 +00:00
self.token_indices = torch.arange(self.image_token_count)
if torch.cuda.is_available():
self.token_indices = self.token_indices.cuda()
def forward(
self,
2022-06-27 15:57:56 +00:00
decoder_state: FloatTensor,
encoder_state: FloatTensor,
attention_state: FloatTensor,
2022-06-27 15:57:56 +00:00
attention_mask: BoolTensor,
token_index: LongTensor
) -> Tuple[FloatTensor, FloatTensor]:
# Self Attention
residual = decoder_state
decoder_state = self.pre_self_attn_layer_norm.forward(decoder_state)
2022-06-29 01:28:36 +00:00
self_attn_mask = self.token_indices < token_index + 1
token_mask = self.token_indices == token_index
2022-06-27 15:57:56 +00:00
self_attn_mask = torch.stack([self_attn_mask] * decoder_state.shape[0])
decoder_state, attention_state = self.self_attn.forward(
2022-06-27 15:57:56 +00:00
decoder_state,
attention_state,
2022-06-27 15:57:56 +00:00
self_attn_mask,
2022-06-29 01:28:36 +00:00
token_mask
2022-06-27 15:57:56 +00:00
)
decoder_state = self.self_attn_layer_norm.forward(decoder_state)
decoder_state = residual + decoder_state
# Cross Attention
residual = decoder_state
decoder_state = self.pre_encoder_attn_layer_norm.forward(decoder_state)
decoder_state = self.encoder_attn.forward(
decoder_state,
encoder_state,
attention_mask
)
decoder_state = self.encoder_attn_layer_norm.forward(decoder_state)
decoder_state = residual + decoder_state
# Feed forward
residual = decoder_state
decoder_state = self.glu.forward(decoder_state)
decoder_state = residual + decoder_state
return decoder_state, attention_state
2022-06-27 15:57:56 +00:00
class DalleBartDecoder(nn.Module):
def __init__(
self,
image_vocab_count: int,
2022-06-27 15:57:56 +00:00
image_token_count: int,
sample_token_count: int,
embed_count: int,
attention_head_count: int,
glu_embed_count: int,
layer_count: int,
batch_count: int,
start_token: int
2022-06-27 15:57:56 +00:00
):
super().__init__()
self.layer_count = layer_count
self.sample_token_count = sample_token_count
2022-06-29 01:28:36 +00:00
self.condition_factor = 10.0
2022-06-27 15:57:56 +00:00
self.image_token_count = image_token_count
self.embed_tokens = nn.Embedding(image_vocab_count + 1, embed_count)
2022-06-27 15:57:56 +00:00
self.embed_positions = nn.Embedding(image_token_count, embed_count)
self.layers: List[DecoderLayer] = nn.ModuleList([
DecoderLayer(
2022-06-27 15:57:56 +00:00
image_token_count,
attention_head_count,
embed_count,
glu_embed_count
)
for _ in range(layer_count)
])
self.layernorm_embedding = nn.LayerNorm(embed_count)
self.final_ln = nn.LayerNorm(embed_count)
self.lm_head = nn.Linear(embed_count, image_vocab_count + 1, bias=False)
self.attention_state_shape = (
layer_count,
2 * batch_count,
2022-06-27 17:19:03 +00:00
image_token_count,
2022-06-29 17:48:12 +00:00
embed_count
2022-06-27 15:57:56 +00:00
)
2022-06-29 01:28:36 +00:00
self.zero_prob = torch.zeros([1])
self.token_indices = torch.arange(self.sample_token_count)
self.start_token = torch.tensor([start_token]).to(torch.long)
if torch.cuda.is_available():
self.zero_prob = self.zero_prob.cuda()
self.token_indices = self.token_indices.cuda()
self.start_token = self.start_token.cuda()
2022-06-27 15:57:56 +00:00
def decode_step(
self,
2022-06-27 15:57:56 +00:00
text_tokens: LongTensor,
encoder_state: FloatTensor,
attention_state: FloatTensor,
prev_token: LongTensor,
token_index: LongTensor
2022-06-27 15:57:56 +00:00
) -> Tuple[LongTensor, FloatTensor]:
2022-06-29 01:28:36 +00:00
attention_mask = text_tokens.not_equal(1)
2022-06-27 15:57:56 +00:00
batch_count = encoder_state.shape[0]
prev_token_batched = torch.cat([prev_token] * batch_count)
token_index_batched = torch.cat([token_index] * batch_count)
decoder_state = self.embed_tokens.forward(prev_token_batched)
decoder_state += self.embed_positions.forward(token_index_batched)
2022-06-27 15:57:56 +00:00
decoder_state = self.layernorm_embedding.forward(decoder_state)
2022-06-27 17:19:03 +00:00
decoder_state = decoder_state[:, None]
attention_states_new = []
for i in range(self.layer_count):
decoder_state, attention_state_layer = self.layers[i].forward(
2022-06-27 15:57:56 +00:00
decoder_state,
encoder_state,
attention_state[i],
2022-06-27 15:57:56 +00:00
attention_mask,
token_index
2022-06-27 15:57:56 +00:00
)
attention_states_new.append(attention_state_layer)
2022-06-27 17:19:03 +00:00
decoder_state = self.final_ln(decoder_state)
logits = self.lm_head(decoder_state)
2022-06-27 15:57:56 +00:00
a = self.condition_factor
2022-07-01 14:17:29 +00:00
logits: FloatTensor = (1 - a) * logits[0, -1] + a * logits[1, -1]
2022-06-27 15:57:56 +00:00
top_logits, _ = logits.topk(50, dim=-1)
2022-06-27 15:57:56 +00:00
probs = torch.where(
logits < top_logits[-1],
2022-06-29 01:28:36 +00:00
self.zero_prob,
2022-06-27 15:57:56 +00:00
torch.exp(logits - top_logits[0])
)
return probs, torch.stack(attention_states_new)
2022-06-27 15:57:56 +00:00
def forward(
self,
2022-06-27 15:57:56 +00:00
text_tokens: LongTensor,
encoder_state: FloatTensor
) -> LongTensor:
image_tokens: List[LongTensor] = []
attention_state = torch.zeros(self.attention_state_shape)
2022-06-29 01:28:36 +00:00
if torch.cuda.is_available():
attention_state = attention_state.cuda()
2022-06-27 15:57:56 +00:00
image_token = self.start_token
for i in range(self.sample_token_count):
probs, attention_state = self.decode_step(
2022-06-27 15:57:56 +00:00
text_tokens = text_tokens,
encoder_state = encoder_state,
attention_state = attention_state,
prev_token = image_token,
token_index = self.token_indices[[i]]
2022-06-27 15:57:56 +00:00
)
image_token = torch.multinomial(probs, 1)
image_tokens += [image_token]
return torch.cat(image_tokens)