Method configure_optimizers

src/model.py:284–345

(self)

            controller.num_att_layers = cross_att_count

    def configure_optimizers(self):
        lr = self.learning_rate
        params = []

        if self.freeze_model == 'crossattn-k':
            for x in self.model.diffusion_model.named_parameters():
                if 'transformer_blocks' in x[0]:
                    if 'attn2.to_k' in x[0]:
                        params += [x[1]]
                        print(x[0])
        elif self.freeze_model == 'crossattn-kv':
            for x in self.model.diffusion_model.named_parameters():
                if 'transformer_blocks' in x[0]:
                    if 'attn2.to_k' in x[0] or 'attn2.to_v' in x[0]:
                        params += [x[1]]
                        print(x[0])
        elif self.freeze_model == 'crossattn-qkv':
            for x in self.model.diffusion_model.named_parameters():
                if 'transformer_blocks' in x[0]:
                    if 'attn2.to_k' in x[0] or 'attn2.to_v' in x[0] or 'attn2.to_q' in x[0]:
                        params += [x[1]]
                        print(x[0])
        elif self.freeze_model == 'selfattn-crossattn-qkv':
            for x in self.model.diffusion_model.named_parameters():
                if 'transformer_blocks' in x[0]:
                    if ('attn2.to_k' in x[0] or 'attn2.to_v' in x[0] or 'attn2.to_q' in x[0] or 'attn1.to_k' in x[0] or 'attn1.to_v' in x[0] or 'attn1.to_q' in x[0]):
                        params += [x[1]]
                        print(x[0])

        elif self.freeze_model == 'crossattn':
            for x in self.model.diffusion_model.named_parameters():
                if 'transformer_blocks' in x[0]:
                    if 'attn2' in x[0]:
                        params += [x[1]]
                        print(x[0])
        else:
            params = list(self.model.parameters())

        if self.cond_stage_trainable:
            print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
            if self.add_token:
                params = params + list(self.cond_stage_model.transformer.text_model.embeddings.token_embedding.parameters())
            else:
                params = params + list(self.cond_stage_model.parameters())

        if self.learn_logvar:
            print('Diffusion model optimizing logvar')
            params.append(self.logvar)
        opt = torch.optim.AdamW(params, lr=lr)
        if self.use_scheduler:
            assert 'target' in self.scheduler_config
            scheduler = instantiate_from_config(self.scheduler_config)

            print("Setting up LambdaLR scheduler...")
            scheduler = [
                {
                    'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
                    'interval': 'step',