File size: 3,252 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import copy
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
import torch
def modulate(x, shift, scale):
if len(x.shape) == 3 and len(shift.shape) == 2:
# [B, K, D] + [B, D]
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
elif len(x.shape) == len(shift.shape) == 3:
# [B, K, D] + [B, K, D]
return x * (1 + scale) + shift
elif len(x.shape) == 4 and len(shift.shape) == 2:
# [B, K, A, D] + [B, D]
return x * (1 + scale.unsqueeze(1).unsqueeze(1)) + shift.unsqueeze(1).unsqueeze(1)
elif len(x.shape) == len(shift.shape) == 4:
# [B, K, A, D] + [B, K, A, D]
return x * (1 + scale) + shift
else:
raise ValueError("Invalid shapes to modulate")
class MTRDecoder(nn.Module):
def __init__(self, config, use_pre_norm, use_adaln=True):
super().__init__()
self.num_blocks = config.get('NUM_DECODER_BLOCKS', 2)
self.self_attn_K = nn.ModuleList([])
self.self_attn_A = nn.ModuleList([])
template_encoder = nn.TransformerEncoderLayer(d_model=config.D_MODEL,
dropout=config.get('DROPOUT_OF_ATTN', 0.1),
nhead=config.NUM_ATTN_HEAD,
dim_feedforward=config.D_MODEL * 4,
norm_first=use_pre_norm,
batch_first=True)
self.use_adaln = use_adaln
if use_adaln:
template_adaln = nn.Sequential(nn.SiLU(),
nn.Linear(config.D_MODEL, 2 * config.D_MODEL, bias=True))
self.t_adaLN = nn.ModuleList([])
for _ in range(self.num_blocks):
self.self_attn_K.append(copy.deepcopy(template_encoder))
self.self_attn_A.append(copy.deepcopy(template_encoder))
if use_adaln:
self.t_adaLN.append(copy.deepcopy(template_adaln))
# zero initialization parameters of adaln
nn.init.constant_(self.t_adaLN[-1][-1].weight, 0)
nn.init.constant_(self.t_adaLN[-1][-1].bias, 0)
def forward(self, query_token, time_emb=None):
"""
@param query_token: [B, K, A, D]
@param time_emb: [B, D]
"""
B, K, A = query_token.shape[:3]
cur_query = query_token
for i in range(self.num_blocks):
if self.use_adaln:
# time modulation
shift, scale = self.t_adaLN[i](time_emb).chunk(2, dim=-1)
cur_query = modulate(cur_query, shift, scale) # [B, K, A, D]
# K-to-K self-attention
cur_query = rearrange(query_token, 'b k a d -> (b a) k d')
cur_query = self.self_attn_K[i](cur_query)
# A-to-A self-attention
cur_query = rearrange(cur_query, '(b a) k d -> (b k) a d', b=B)
cur_query = self.self_attn_A[i](cur_query)
# reshape
cur_query = rearrange(cur_query, '(b k) a d -> b k a d', b=B)
return cur_query
|