Method state_dict

networks/dylora.py:123–149

(self, destination=None, prefix="", keep_vars=False)

        return result

    def state_dict(self, destination=None, prefix="", keep_vars=False):
        # state dictを通常のLoRAと同じにする:
        # nn.ParameterListは `.lora_A.0` みたいな名前になるので、forwardと同様にcatして入れ替える
        sd = super().state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)

        lora_A_weight = torch.cat(tuple(self.lora_A), dim=0)
        if self.is_conv2d and not self.is_conv2d_3x3:
            lora_A_weight = lora_A_weight.unsqueeze(-1).unsqueeze(-1)

        lora_B_weight = torch.cat(tuple(self.lora_B), dim=1)
        if self.is_conv2d and not self.is_conv2d_3x3:
            lora_B_weight = lora_B_weight.unsqueeze(-1).unsqueeze(-1)

        sd[self.lora_name + ".lora_down.weight"] = lora_A_weight if keep_vars else lora_A_weight.detach()
        sd[self.lora_name + ".lora_up.weight"] = lora_B_weight if keep_vars else lora_B_weight.detach()

        i = 0
        while True:
            key_a = f"{self.lora_name}.lora_A.{i}"
            key_b = f"{self.lora_name}.lora_B.{i}"
            if key_a in sd:
                sd.pop(key_a)
                sd.pop(key_b)
            else:
                break
            i += 1
        return sd

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        # 通常のLoRAと同じstate dictを読み込めるようにする：この方法はchatGPTに聞いた