diff options
Diffstat (limited to 'training')
| -rw-r--r-- | training/optimization.py | 2 | ||||
| -rw-r--r-- | training/ti.py | 3 |
2 files changed, 3 insertions, 2 deletions
diff --git a/training/optimization.py b/training/optimization.py index c501ed9..3809f3b 100644 --- a/training/optimization.py +++ b/training/optimization.py | |||
| @@ -6,7 +6,7 @@ from diffusers.utils import logging | |||
| 6 | logger = logging.get_logger(__name__) | 6 | logger = logging.get_logger(__name__) |
| 7 | 7 | ||
| 8 | 8 | ||
| 9 | def get_one_cycle_schedule(optimizer, num_training_steps, annealing="cos", min_lr=0.001, mid_point=0.4, last_epoch=-1): | 9 | def get_one_cycle_schedule(optimizer, num_training_steps, annealing="cos", min_lr=0.01, mid_point=0.4, last_epoch=-1): |
| 10 | """ | 10 | """ |
| 11 | Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after | 11 | Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after |
| 12 | a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. | 12 | a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. |
diff --git a/training/ti.py b/training/ti.py index 1318e22..031fe48 100644 --- a/training/ti.py +++ b/training/ti.py | |||
| @@ -22,7 +22,8 @@ class TrainableEmbeddings(CLIPTextEmbeddings): | |||
| 22 | self.train_indices = torch.tensor(new_ids) | 22 | self.train_indices = torch.tensor(new_ids) |
| 23 | 23 | ||
| 24 | self.trainable_embedding = nn.Embedding(self.token_embedding.num_embeddings, self.token_embedding.embedding_dim) | 24 | self.trainable_embedding = nn.Embedding(self.token_embedding.num_embeddings, self.token_embedding.embedding_dim) |
| 25 | self.trainable_embedding.weight.data = self.token_embedding.weight.data.clone() | 25 | self.trainable_embedding.weight.data.zero_() |
| 26 | self.trainable_embedding.weight.data[self.train_indices] = self.token_embedding.weight.data[self.train_indices] | ||
| 26 | 27 | ||
| 27 | def forward( | 28 | def forward( |
| 28 | self, | 29 | self, |