From 8364ce697ddf6117fdd4f7222832d546d63880de Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Wed, 21 Jun 2023 13:28:49 +0200
Subject: Update
---
models/attention/control.py | 106 ++++++++++++++++++++++++++++-----------
models/attention/hook.py | 5 +-
models/attention/structured.py | 65 ++++++++++++++----------
models/clip/embeddings.py | 29 ++++++-----
models/clip/tokenizer.py | 23 +++++----
models/clip/util.py | 17 +++++--
models/convnext/discriminator.py | 11 ++--
models/sparse.py | 12 +++--
8 files changed, 176 insertions(+), 92 deletions(-)
(limited to 'models')
diff --git a/models/attention/control.py b/models/attention/control.py
index 248bd9f..ec378c4 100644
--- a/models/attention/control.py
+++ b/models/attention/control.py
@@ -23,7 +23,7 @@ class AttentionControl(abc.ABC):
attn = self.forward(attn, is_cross, place_in_unet)
else:
h = attn.shape[0]
- attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
+ attn[h // 2 :] = self.forward(attn[h // 2 :], is_cross, place_in_unet)
self.cur_att_layer += 1
if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
self.cur_att_layer = 0
@@ -49,12 +49,18 @@ class EmptyControl(AttentionControl):
class AttentionStore(AttentionControl):
@staticmethod
def get_empty_store():
- return {"down_cross": [], "mid_cross": [], "up_cross": [],
- "down_self": [], "mid_self": [], "up_self": []}
+ return {
+ "down_cross": [],
+ "mid_cross": [],
+ "up_cross": [],
+ "down_self": [],
+ "mid_self": [],
+ "up_self": [],
+ }
def forward(self, attn, is_cross: bool, place_in_unet: str):
key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
- if attn.shape[1] <= 32 ** 2: # avoid memory overhead
+ if attn.shape[1] <= 32**2: # avoid memory overhead
self.step_store[key].append(attn)
return attn
@@ -68,8 +74,10 @@ class AttentionStore(AttentionControl):
self.step_store = self.get_empty_store()
def get_average_attention(self):
- average_attention = {key: [item / self.cur_step for item in self.attention_store[key]]
- for key in self.attention_store}
+ average_attention = {
+ key: [item / self.cur_step for item in self.attention_store[key]]
+ for key in self.attention_store
+ }
return average_attention
def reset(self):
@@ -90,7 +98,7 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
return x_t
def replace_self_attention(self, attn_base, att_replace):
- if att_replace.shape[2] <= 16 ** 2:
+ if att_replace.shape[2] <= 16**2:
return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
else:
return att_replace
@@ -101,41 +109,62 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
def forward(self, attn, is_cross: bool, place_in_unet: str):
super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
- if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):
+ if is_cross or (
+ self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]
+ ):
h = attn.shape[0] // (self.batch_size)
attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
attn_base, attn_repalce = attn[0], attn[1:]
if is_cross:
alpha_words = self.cross_replace_alpha[self.cur_step]
- attn_repalce_new = self.replace_cross_attention(
- attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce
+ attn_repalce_new = (
+ self.replace_cross_attention(attn_base, attn_repalce) * alpha_words
+ + (1 - alpha_words) * attn_repalce
+ )
attn[1:] = attn_repalce_new
else:
attn[1:] = self.replace_self_attention(attn_base, attn_repalce)
attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
return attn
- def __init__(self, prompts, num_steps: int,
- cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
- self_replace_steps: Union[float, Tuple[float, float]],
- local_blend: Optional[LocalBlend]):
+ def __init__(
+ self,
+ prompts,
+ num_steps: int,
+ cross_replace_steps: Union[
+ float, Tuple[float, float], Dict[str, Tuple[float, float]]
+ ],
+ self_replace_steps: Union[float, Tuple[float, float]],
+ local_blend: Optional[LocalBlend],
+ ):
super(AttentionControlEdit, self).__init__()
self.batch_size = len(prompts)
self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(
- prompts, num_steps, cross_replace_steps, tokenizer).to(device)
+ prompts, num_steps, cross_replace_steps, tokenizer
+ ).to(device)
if type(self_replace_steps) is float:
self_replace_steps = 0, self_replace_steps
- self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])
+ self.num_self_replace = int(num_steps * self_replace_steps[0]), int(
+ num_steps * self_replace_steps[1]
+ )
self.local_blend = local_blend
class AttentionReplace(AttentionControlEdit):
def replace_cross_attention(self, attn_base, att_replace):
- return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)
-
- def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
- local_blend: Optional[LocalBlend] = None):
- super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
+ return torch.einsum("hpw,bwn->bhpn", attn_base, self.mapper)
+
+ def __init__(
+ self,
+ prompts,
+ num_steps: int,
+ cross_replace_steps: float,
+ self_replace_steps: float,
+ local_blend: Optional[LocalBlend] = None,
+ ):
+ super(AttentionReplace, self).__init__(
+ prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
+ )
self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device)
@@ -145,9 +174,17 @@ class AttentionRefine(AttentionControlEdit):
attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
return attn_replace
- def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
- local_blend: Optional[LocalBlend] = None):
- super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
+ def __init__(
+ self,
+ prompts,
+ num_steps: int,
+ cross_replace_steps: float,
+ self_replace_steps: float,
+ local_blend: Optional[LocalBlend] = None,
+ ):
+ super(AttentionRefine, self).__init__(
+ prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
+ )
self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)
self.mapper, alphas = self.mapper.to(device), alphas.to(device)
self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])
@@ -156,13 +193,24 @@ class AttentionRefine(AttentionControlEdit):
class AttentionReweight(AttentionControlEdit):
def replace_cross_attention(self, attn_base, att_replace):
if self.prev_controller is not None:
- attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace)
+ attn_base = self.prev_controller.replace_cross_attention(
+ attn_base, att_replace
+ )
attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]
return attn_replace
- def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer,
- local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None):
- super(AttentionReweight, self).__init__(prompts, num_steps,
- cross_replace_steps, self_replace_steps, local_blend)
+ def __init__(
+ self,
+ prompts,
+ num_steps: int,
+ cross_replace_steps: float,
+ self_replace_steps: float,
+ equalizer,
+ local_blend: Optional[LocalBlend] = None,
+ controller: Optional[AttentionControlEdit] = None,
+ ):
+ super(AttentionReweight, self).__init__(
+ prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
+ )
self.equalizer = equalizer.to(device)
self.prev_controller = controller
diff --git a/models/attention/hook.py b/models/attention/hook.py
index 903de02..6b5fb68 100644
--- a/models/attention/hook.py
+++ b/models/attention/hook.py
@@ -3,6 +3,7 @@ import torch
try:
import xformers.ops
+
xformers._is_functorch_available = True
MEM_EFFICIENT_ATTN = True
except ImportError:
@@ -42,10 +43,10 @@ def register_attention_control(model, controller):
return forward
def register_recr(net_, count, place_in_unet):
- if net_.__class__.__name__ == 'CrossAttention':
+ if net_.__class__.__name__ == "CrossAttention":
net_.forward = ca_forward(net_, place_in_unet)
return count + 1
- elif hasattr(net_, 'children'):
+ elif hasattr(net_, "children"):
for net__ in net_.children():
count = register_recr(net__, count, place_in_unet)
return count
diff --git a/models/attention/structured.py b/models/attention/structured.py
index 24d889f..5bbbc06 100644
--- a/models/attention/structured.py
+++ b/models/attention/structured.py
@@ -16,7 +16,9 @@ class StructuredAttentionControl(AttentionControl):
if self.struct_attn:
out = self.struct_qkv(q, context, mask)
else:
- context = torch.cat([context[0], context[1]['k'][0]], dim=0) # use key tensor for context
+ context = torch.cat(
+ [context[0], context[1]["k"][0]], dim=0
+ ) # use key tensor for context
out = self.normal_qkv(q, context, mask)
else:
context = default(context, x)
@@ -29,11 +31,13 @@ class StructuredAttentionControl(AttentionControl):
context: list of [uc, list of conditional context]
"""
uc_context = context[0]
- context_k, context_v = context[1]['k'], context[1]['v']
+ context_k, context_v = context[1]["k"], context[1]["v"]
if isinstance(context_k, list) and isinstance(context_v, list):
out = self.multi_qkv(q, uc_context, context_k, context_v, mask)
- elif isinstance(context_k, torch.Tensor) and isinstance(context_v, torch.Tensor):
+ elif isinstance(context_k, torch.Tensor) and isinstance(
+ context_v, torch.Tensor
+ ):
out = self.heterogeous_qkv(q, uc_context, context_k, context_v, mask)
else:
raise NotImplementedError
@@ -50,36 +54,45 @@ class StructuredAttentionControl(AttentionControl):
k_c = [self.to_k(c_k) for c_k in context_k]
v_c = [self.to_v(c_v) for c_v in context_v]
- q = rearrange(q, 'b n (h d) -> (b h) n d', h=h)
+ q = rearrange(q, "b n (h d) -> (b h) n d", h=h)
- k_uc = rearrange(k_uc, 'b n (h d) -> (b h) n d', h=h)
- v_uc = rearrange(v_uc, 'b n (h d) -> (b h) n d', h=h)
+ k_uc = rearrange(k_uc, "b n (h d) -> (b h) n d", h=h)
+ v_uc = rearrange(v_uc, "b n (h d) -> (b h) n d", h=h)
- k_c = [rearrange(k, 'b n (h d) -> (b h) n d', h=h) for k in k_c] # NOTE: modification point
- v_c = [rearrange(v, 'b n (h d) -> (b h) n d', h=h) for v in v_c]
+ k_c = [
+ rearrange(k, "b n (h d) -> (b h) n d", h=h) for k in k_c
+ ] # NOTE: modification point
+ v_c = [rearrange(v, "b n (h d) -> (b h) n d", h=h) for v in v_c]
# get composition
- sim_uc = einsum('b i d, b j d -> b i j', q[:true_bs], k_uc) * self.scale
- sim_c = [einsum('b i d, b j d -> b i j', q[true_bs:], k) * self.scale for k in k_c]
+ sim_uc = einsum("b i d, b j d -> b i j", q[:true_bs], k_uc) * self.scale
+ sim_c = [
+ einsum("b i d, b j d -> b i j", q[true_bs:], k) * self.scale for k in k_c
+ ]
attn_uc = sim_uc.softmax(dim=-1)
attn_c = [sim.softmax(dim=-1) for sim in sim_c]
# get uc output
- out_uc = einsum('b i j, b j d -> b i d', attn_uc, v_uc)
+ out_uc = einsum("b i j, b j d -> b i d", attn_uc, v_uc)
# get c output
if len(v_c) == 1:
out_c_collect = []
for attn in attn_c:
for v in v_c:
- out_c_collect.append(einsum('b i j, b j d -> b i d', attn, v))
+ out_c_collect.append(einsum("b i j, b j d -> b i d", attn, v))
out_c = sum(out_c_collect) / len(out_c_collect)
else:
- out_c = sum([einsum('b i j, b j d -> b i d', attn, v) for attn, v in zip(attn_c, v_c)]) / len(v_c)
+ out_c = sum(
+ [
+ einsum("b i j, b j d -> b i d", attn, v)
+ for attn, v in zip(attn_c, v_c)
+ ]
+ ) / len(v_c)
out = torch.cat([out_uc, out_c], dim=0)
- out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+ out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
return out
@@ -88,21 +101,21 @@ class StructuredAttentionControl(AttentionControl):
k = self.to_k(context)
v = self.to_v(context)
- q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+ q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
- sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+ sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
if exists(mask):
- mask = rearrange(mask, 'b ... -> b (...)')
+ mask = rearrange(mask, "b ... -> b (...)")
max_neg_value = -torch.finfo(sim.dtype).max
- mask = repeat(mask, 'b j -> (b h) () j', h=h)
+ mask = repeat(mask, "b j -> (b h) () j", h=h)
sim.masked_fill_(~mask, max_neg_value)
# attention, what we cannot get enough of
attn = sim.softmax(dim=-1)
- out = einsum('b i j, b j d -> b i d', attn, v)
- out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+ out = einsum("b i j, b j d -> b i d", attn, v)
+ out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
return out
@@ -111,21 +124,21 @@ class StructuredAttentionControl(AttentionControl):
k = self.to_k(torch.cat([uc_context, context_k], dim=0))
v = self.to_v(torch.cat([uc_context, context_v], dim=0))
- q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+ q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
- sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+ sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
if exists(mask):
- mask = rearrange(mask, 'b ... -> b (...)')
+ mask = rearrange(mask, "b ... -> b (...)")
max_neg_value = -torch.finfo(sim.dtype).max
- mask = repeat(mask, 'b j -> (b h) () j', h=h)
+ mask = repeat(mask, "b j -> (b h) () j", h=h)
sim.masked_fill_(~mask, max_neg_value)
# attention, what we cannot get enough of
attn = sim.softmax(dim=-1)
- out = einsum('b i j, b j d -> b i d', attn, v)
- out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+ out = einsum("b i j, b j d -> b i d", attn, v)
+ out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
return out
def get_kv(self, context):
diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 7c7f2ac..8c3c6d4 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -14,7 +14,13 @@ from models.sparse import SparseEmbedding
class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
- def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, alpha: int = 8, dropout: float = 0.0):
+ def __init__(
+ self,
+ config: CLIPTextConfig,
+ embeddings: CLIPTextEmbeddings,
+ alpha: int = 8,
+ dropout: float = 0.0,
+ ):
super().__init__(config)
self.position_embedding = embeddings.position_embedding
@@ -28,7 +34,9 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
self.token_embedding.weight = embeddings.token_embedding.weight
def resize(self, size: int):
- self.token_embedding = self.token_embedding.new_resized(size, self.initializer_factor)
+ self.token_embedding = self.token_embedding.new_resized(
+ size, self.initializer_factor
+ )
def add_embed(
self,
@@ -46,7 +54,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
initializer = [initializer]
if isinstance(initializer, list):
- initializer = (initializer * len(token_ids))[:len(token_ids)]
+ initializer = (initializer * len(token_ids))[: len(token_ids)]
with torch.no_grad():
initializer = self.get_embed(initializer)
@@ -76,24 +84,21 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
def get_embed(self, input_ids: Union[list[int], torch.LongTensor]):
if isinstance(input_ids, list):
- input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device, dtype=torch.long)
+ input_ids = torch.tensor(
+ input_ids, device=self.token_embedding.weight.device, dtype=torch.long
+ )
return self.token_embedding(input_ids)
def patch_managed_embeddings(
- text_encoder: CLIPTextModel,
- alpha: int = 8,
- dropout: float = 0.0
+ text_encoder: CLIPTextModel, alpha: int = 8, dropout: float = 0.0
) -> ManagedCLIPTextEmbeddings:
if isinstance(text_encoder.text_model.embeddings, ManagedCLIPTextEmbeddings):
return text_encoder.text_model.embeddings
-
+
text_embeddings = ManagedCLIPTextEmbeddings(
- text_encoder.config,
- text_encoder.text_model.embeddings,
- alpha,
- dropout
+ text_encoder.config, text_encoder.text_model.embeddings, alpha, dropout
)
text_encoder.text_model.embeddings = text_embeddings
return text_embeddings
diff --git a/models/clip/tokenizer.py b/models/clip/tokenizer.py
index 789b525..a866641 100644
--- a/models/clip/tokenizer.py
+++ b/models/clip/tokenizer.py
@@ -91,18 +91,21 @@ class MultiCLIPTokenizer(CLIPTokenizer):
self.vector_shuffle = shuffle_none
def add_multi_tokens(
- self,
- new_tokens: Union[str, list[str]],
- num_vectors: Union[int, list[int]] = 1
+ self, new_tokens: Union[str, list[str]], num_vectors: Union[int, list[int]] = 1
) -> Union[list[int], list[list[int]]]:
if isinstance(new_tokens, list):
if isinstance(num_vectors, int):
num_vectors = [num_vectors] * len(new_tokens)
if len(num_vectors) != len(new_tokens):
- raise ValueError("Expected new_tokens and num_vectors to have the same len")
+ raise ValueError(
+ "Expected new_tokens and num_vectors to have the same len"
+ )
- return [self.add_multi_tokens(new_token, vecs) for new_token, vecs in zip(new_tokens, num_vectors)]
+ return [
+ self.add_multi_tokens(new_token, vecs)
+ for new_token, vecs in zip(new_tokens, num_vectors)
+ ]
if isinstance(num_vectors, list):
raise ValueError("Expected num_vectors to be int for single token")
@@ -129,13 +132,11 @@ class MultiCLIPTokenizer(CLIPTokenizer):
return [id]
def expand_ids(self, ids: list[int]):
- return [
- new_id
- for id in ids
- for new_id in self.expand_id(id)
- ]
+ return [new_id for id in ids for new_id in self.expand_id(id)]
- def expand_batched_ids(self, input_ids: Union[list[int], list[list[int]], tuple[list[int]]]):
+ def expand_batched_ids(
+ self, input_ids: Union[list[int], list[list[int]], tuple[list[int]]]
+ ):
if isinstance(input_ids, (list, tuple)) and isinstance(input_ids[0], list):
return [self.expand_ids(batch) for batch in input_ids]
else:
diff --git a/models/clip/util.py b/models/clip/util.py
index f94fbc7..7196bb6 100644
--- a/models/clip/util.py
+++ b/models/clip/util.py
@@ -5,27 +5,32 @@ import torch
from transformers import CLIPTokenizer, CLIPTextModel
-def unify_input_ids(tokenizer: CLIPTokenizer, input_ids: list[list[int]], max_length: Optional[int] = None):
+def unify_input_ids(
+ tokenizer: CLIPTokenizer,
+ input_ids: list[list[int]],
+ max_length: Optional[int] = None,
+):
if max_length is None:
return tokenizer.pad(
{"input_ids": input_ids},
padding=True,
pad_to_multiple_of=tokenizer.model_max_length,
- return_tensors="pt"
+ return_tensors="pt",
)
else:
return tokenizer.pad(
{"input_ids": input_ids},
padding="max_length",
max_length=max_length,
- return_tensors="pt"
+ return_tensors="pt",
)
+
def get_extended_embeddings(
text_encoder: CLIPTextModel,
input_ids: torch.LongTensor,
position_ids: Optional[torch.LongTensor] = None,
- attention_mask=None
+ attention_mask=None,
):
model_max_length = text_encoder.config.max_position_embeddings
prompts = input_ids.shape[0]
@@ -36,6 +41,8 @@ def get_extended_embeddings(
if attention_mask is not None:
attention_mask = attention_mask.view((-1, model_max_length))
- text_embeddings = text_encoder(input_ids, position_ids=position_ids, attention_mask=attention_mask)[0]
+ text_embeddings = text_encoder(
+ input_ids, position_ids=position_ids, attention_mask=attention_mask
+ )[0]
text_embeddings = text_embeddings.view((prompts, -1, text_embeddings.shape[2]))
return text_embeddings
diff --git a/models/convnext/discriminator.py b/models/convnext/discriminator.py
index 571b915..5798bcf 100644
--- a/models/convnext/discriminator.py
+++ b/models/convnext/discriminator.py
@@ -5,7 +5,7 @@ from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from torch.nn import functional as F
-class ConvNeXtDiscriminator():
+class ConvNeXtDiscriminator:
def __init__(self, model: ConvNeXt, input_size: int) -> None:
self.net = model
@@ -22,8 +22,13 @@ class ConvNeXtDiscriminator():
img_mean = self.img_mean.to(device=img.device, dtype=img.dtype)
img_std = self.img_std.to(device=img.device, dtype=img.dtype)
- img = ((img + 1.) / 2.).sub(img_mean).div(img_std)
+ img = ((img + 1.0) / 2.0).sub(img_mean).div(img_std)
- img = F.interpolate(img, size=(self.input_size, self.input_size), mode='bicubic', align_corners=True)
+ img = F.interpolate(
+ img,
+ size=(self.input_size, self.input_size),
+ mode="bicubic",
+ align_corners=True,
+ )
pred = self.net(img)
return pred
diff --git a/models/sparse.py b/models/sparse.py
index bd45696..e5897c9 100644
--- a/models/sparse.py
+++ b/models/sparse.py
@@ -15,21 +15,25 @@ class SparseEmbedding(nn.Embedding):
):
nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
- self.register_buffer('trainable_ids', self.weight.new_zeros(num_embeddings, dtype=torch.long) - 1)
+ self.register_buffer(
+ "trainable_ids", self.weight.new_zeros(num_embeddings, dtype=torch.long) - 1
+ )
self.trainable = nn.ParameterList()
self.scaling = alpha
self.dropout_p = dropout
self.weight.requires_grad = False
- if dropout > 0.:
+ if dropout > 0.0:
self.dropout = nn.Dropout(p=dropout)
else:
self.dropout = nn.Identity()
self.reset_parameters()
- def new_resized(self, new_num_embeddings: int, initializer_factor: Optional[float] = None):
+ def new_resized(
+ self, new_num_embeddings: int, initializer_factor: Optional[float] = None
+ ):
n = min(self.num_embeddings, new_num_embeddings)
new_emb = SparseEmbedding(
@@ -38,7 +42,7 @@ class SparseEmbedding(nn.Embedding):
self.scaling,
self.dropout_p,
device=self.weight.device,
- dtype=self.weight.dtype
+ dtype=self.weight.dtype,
)
if initializer_factor is not None:
new_emb.weight.data.normal_(mean=0.0, std=initializer_factor * 0.02)
--
cgit v1.2.3-70-g09d2