我想要一天分享一點「LLM從底層堆疊的技術」，並且每篇文章長度控制在三分鐘以內，讓大家不會壓力太大，但是又能夠每天成長一點。
Keras 實現 Stable Diffusion 的 Down-Sampling 部分為：
class DiffusionModelV2(keras.Model):
def __init__(self, img_height, img_width, max_text_length, name = None, download_weights = True):
context = keras.layers.Input((max_text_length, 1024))
t_embed_input = keras.layers.Input((320,))
latent = keras.layers.Input((img_height // 8, img_width // 8, 4))
t_emb = keras.layers.Dense(1280)(t_embed_input)
t_emb = keras.layers.Activation("swish")(t_emb)
t_emb = keras.layers.Dense(1280)(t_emb)
# Downsampling Flow
outputs = []
x = PaddedConv2D(320, kernel_size = 3, padding = 1)(latent)
outputs.append(x)
for _ in range(2):
x = ResBlock(320)([x, t_emb])
x = SpatialTransformer(5, 64, fully_connected = True)([x, context])
outputs.append(x)
x = PaddedConv2D(320, 3, strides = 2, padding = 1)(x) # Downsample 2x
outputs.append(x)
for _ in range(2):
x = ResBlock(640)([x, t_emb])
x = SpatialTransformer(10, 64, fully_connected = True)([x, context])
outputs.append(x)
x = PaddedConv2D(640, 3, strides = 2, padding = 1)(x) # Downsample 2x
outputs.append(x)
for _ in range(2):
x = ResBlock(1280)([x, t_emb])
x = SpatialTransformer(20, 64, fully_connected = True)([x, context])
outputs.append(x)
x = PaddedConv2D(1280, 3, strides = 2, padding = 1)(x) # Downsample 2x
outputs.append(x)
for _ in range(2):
x = ResBlock(1280)([x, t_emb])
outputs.append(x)
# Middle Flow
x = ResBlock(1280)([x, t_emb])
x = SpatialTransformer(20, 64, fully_connected = True)([x, context])
x = ResBlock(1280)([x, t_emb])
# Upsampling Flow
for _ in range(3):
x = keras.layers.Concatenate()([x, outputs.pop()])
x = ResBlock(1280)([x, t_emb])
x = Upsample(1280)(x)
for _ in range(3):
x = keras.layers.Concatenate()([x, outputs.pop()])
x = ResBlock(1280)([x, t_emb])
x = SpatialTransformer(20, 64, fully_connected = True)([x, context])
x = Upsample(1280)(x)
for _ in range(3):
x = keras.layers.Concatenate()([x, outputs.pop()])
x = ResBlock(640)([x, t_emb])
x = SpatialTransformer(10, 64, fully_connected = True)([x, context])
x = Upsample(640)(x)
for _ in range(3):
x = keras.layers.Concatenate()([x, outputs.pop()])
x = ResBlock(320)([x, t_emb])
x = SpatialTransformer(5, 64, fully_connected = True)([x, context])
# Exit Flow
x = keras.layers.GroupNormalization(epsilon = 1e-5)(x)
x = keras.layers.Activation("swish")(x)
output = PaddedConv2D(4, kernel_size = 3, padding = 1)(x)
當中關鍵區塊程式的源碼為：
class ResBlock(keras.layers.Layer):
def __init__(self, output_dim, **kwargs):
super().__init__(**kwargs)
self.output_dim = output_dim
self.entry_flow = [keras.layers.GroupNormalization(epsilon = 1e-5),
keras.layers.Activation("swish"),
PaddedConv2D(output_dim, 3, padding = 1),]
self.embedding_flow = [keras.layers.Activation("swish"),
keras.layers.Dense(output_dim),]
self.exit_flow = [keras.layers.GroupNormalization(epsilon = 1e-5),
keras.layers.Activation("swish"),
PaddedConv2D(output_dim, 3, padding = 1),]
def call(self, inputs):
inputs, embeddings = inputs
x = inputs
for layer in self.entry_flow:
x = layer(x)
for layer in self.embedding_flow:
embeddings = layer(embeddings)
x = x + embeddings[:, None, None] # Channel Dim is 3 in Keras
for layer in self.exit_flow:
x = layer(x)
return x + self.residual_projection(inputs)
class SpatialTransformer(keras.layers.Layer):
def __init__(self, num_heads, head_size, fully_connected = False, **kwargs):
super().__init__(**kwargs)
self.norm = keras.layers.GroupNormalization(epsilon = 1e-5)
channels = num_heads * head_size
if fully_connected:
self.proj1 = keras.layers.Dense(num_heads * head_size)
else:
self.proj1 = PaddedConv2D(num_heads * head_size, 1)
self.transformer_block = BasicTransformerBlock(channels, num_heads, head_size)
if fully_connected:
self.proj2 = keras.layers.Dense(channels)
else:
self.proj2 = PaddedConv2D(channels, 1)
def call(self, inputs):
inputs, context = inputs
_, h, w, c = inputs.shape
x = self.norm(inputs)
x = self.proj1(x)
x = ops.reshape(x, (-1, h * w, c))
x = self.transformer_block([x, context]) # 注意 x 的最後維度和 context 最後維度可以不同，因為 Attention 機制一進去會先做 Dense
x = ops.reshape(x, (-1, h, w, c))
return self.proj2(x) + inputs
class BasicTransformerBlock(keras.layers.Layer):
def __init__(self, dim, num_heads, head_size, **kwargs):
super().__init__(**kwargs)
self.norm1 = keras.layers.LayerNormalization(epsilon = 1e-5)
self.attn1 = CrossAttention(num_heads, head_size)
self.norm2 = keras.layers.LayerNormalization(epsilon = 1e-5)
self.attn2 = CrossAttention(num_heads, head_size)
self.norm3 = keras.layers.LayerNormalization(epsilon = 1e-5)
self.geglu = GEGLU(dim * 4)
self.dense = keras.layers.Dense(dim)
def call(self, inputs):
inputs, context = inputs
x = self.attn1(self.norm1(inputs), context = None) + inputs
x = self.attn2(self.norm2(x), context = context) + x
return self.dense(self.geglu(self.norm3(x))) + x
class CrossAttention(keras.layers.Layer):
def __init__(self, num_heads, head_size, **kwargs):
super().__init__(**kwargs)
self.to_q = keras.layers.Dense(num_heads * head_size, use_bias = False)
self.to_k = keras.layers.Dense(num_heads * head_size, use_bias = False)
self.to_v = keras.layers.Dense(num_heads * head_size, use_bias = False)
self.scale = head_size**-0.5
self.num_heads = num_heads
self.head_size = head_size
self.out_proj = keras.layers.Dense(num_heads * head_size)
def call(self, inputs, context = None):
if context is None:
context = inputs
q, k, v = self.to_q(inputs), self.to_k(context), self.to_v(context)
q = ops.reshape(q, (-1, inputs.shape[1], self.num_heads, self.head_size))
k = ops.reshape(k, (-1, context.shape[1], self.num_heads, self.head_size))
v = ops.reshape(v, (-1, context.shape[1], self.num_heads, self.head_size))
q = ops.transpose(q, (0, 2, 1, 3)) # (bs, num_heads, time, head_size) # 注意這裡的 time 是圖片的 h x w
k = ops.transpose(k, (0, 2, 3, 1)) # (bs, num_heads, head_size, time) # 注意這裡的 time 是 context 長 ; 不一定等於 h x w
v = ops.transpose(v, (0, 2, 1, 3)) # (bs, num_heads, time, head_size) # 注意這裡的 time 是 context 長 ; 不一定等於 h x w
score = td_dot(q, k) * self.scale
weights = keras.activations.softmax(score) # (bs, num_heads, time, time)
attn = td_dot(weights, v)
attn = ops.transpose(attn, (0, 2, 1, 3)) # (bs, time, num_heads, head_size)
out = ops.reshape(attn, (-1, inputs.shape[1], self.num_heads * self.head_size))
return self.out_proj(out)
解析如下：
- embeddings[:, None, None]
- 形狀從 (batch_size, embedding_dim) 變為 (batch_size, 1, 1, embedding_dim)
- 第一個 None 在原來的第 1 維插入一個新維度
- 第二個 None 在原來的第 2 維插入一個新維度
