Nagi-ovo

import torch
import numpy as np
_ = torch.manual_seed(0)

d, k = 10, 10
W_rank = 2
W = torch.randn(d,W_rank) @ torch.randn(W_rank,k)

W_rank = np.linalg.matrix_rank(W) print(f'Rank of W: {W_rank}')
print(f"{W_rank=}")

# W に対して SVD を実行する (W = UxSxV^T)
U, S, V = torch.svd(W)

# ランク-r の因子分解のため、最初の r の特異値（および U と V の対応する列）だけを保持します
U_r = U[:, :W_rank]
S_r = torch.diag(S[:W_rank])
V_r = V[:, :W_rank].t()  # V_r を転置して正しい次元を得ます

# B = U_r * S_r と A = V_r を計算します
B = U_r @ S_r
A = V_r
print(f'Shape of B: {B.shape}')
print(f'Shape of A: {A.shape}')

U_r = U[:, :W_rank]
S_r = torch.diag(S[:W_rank]) # 特異値の対角行列を得ます
V_r = V[:, :W_rank].t()

B = U_r @ S_r
A = V_r

# ランダムなバイアスと入力を生成します
bias = torch.randn(d)
x = torch.randn(d)

# y = Wx + bias を計算します
y = W @ x + bias

# y' = (B*A)x + bias を計算します
y_prime = (B @ A) @ x + bias

# 二つの結果がほぼ等しいか確認します
if torch.allclose(y, y_prime, rtol=1e-05, atol=1e-08):
    print("y と y' はほぼ等しいです。")
else:
    print("y と y' は等しくありません。")

# MNIST の数字を分類するために過剰に高価なニューラルネットワークを作成します
# お金持ちのパパがいるので、効率は気にしません
class RichBoyNet(nn.Module):
    def __init__(self, hidden_size_1=1000, hidden_size_2=2000):
        super(RichBoyNet,self).__init__()
        self.linear1 = nn.Linear(28*28, hidden_size_1) 
        self.linear2 = nn.Linear(hidden_size_1, hidden_size_2) 
        self.linear3 = nn.Linear(hidden_size_2, 10)
        self.relu = nn.ReLU()

    def forward(self, img):
        x = img.view(-1, 28*28)
        x = self.relu(self.linear1(x))
        x = self.relu(self.linear2(x))
        x = self.linear3(x)
        return x

net = RichBoyNet().to(device)

train(train_loader, net, epochs=1)

class LoRAParametrization(nn.Module):
    def __init__(self, features_in, features_out, rank=1, alpha=1, device='cpu'):
        super().__init__()
        # Aをガウス分布で初期化し、Bをゼロで初期化して、訓練開始時に∆W = BAがゼロになるようにします
        self.lora_A = nn.Parameter(torch.zeros((rank,features_out)).to(device))
        self.lora_B = nn.Parameter(torch.zeros((features_in, rank)).to(device))
        nn.init.normal_(self.lora_A, mean=0, std=1)
        
        # 論文4.1に基づき、スケール因子α/rでハイパーパラメータの調整を簡略化し、αを最初に試すr値に設定します
        self.scale = alpha / rank
        self.enabled = True

    def forward(self, original_weights):
        if self.enabled:
            # W + (B*A)*scale を返します
            return original_weights + torch.matmul(self.lora_B, self.lora_A).view(original_weights.shape) * self.scale
        else:
            return original_weights

import torch.nn.utils.parametrize as parametrize

def linear_layer_parameterization(layer, device, rank=1, lora_alpha=1):
    # バイアスを無視して、重み行列にのみパラメータ化を追加します
    features_in, features_out = layer.weight.shape
    return LoRAParametrization(
        features_in, features_out, rank=rank, alpha=lora_alpha, device=device
    )

parametrize.register_parametrization(
    net.linear1, "weight", linear_layer_parameterization(net.linear1, device)
)
parametrize.register_parametrization(
    net.linear2, "weight", linear_layer_parameterization(net.linear2, device)
)
parametrize.register_parametrization(
    net.linear3, "weight", linear_layer_parameterization(net.linear3, device)
)

def enable_disable_lora(enabled=True):
    for layer in [net.linear1, net.linear2, net.linear3]:
        layer.parametrizations["weight"][0].enabled = enabled

total_parameters_lora = 0
total_parameters_non_lora = 0
for index, layer in enumerate([net.linear1, net.linear2, net.linear3]):
    total_parameters_lora += layer.parametrizations["weight"][0].lora_A.nelement() + layer.parametrizations["weight"][0].lora_B.nelement()
    total_parameters_non_lora += layer.weight.nelement() + layer.bias.nelement()
    print(
        f'Layer {index+1}: W: {layer.weight.shape} + B: {layer.bias.shape} + Lora_A: {layer.parametrizations["weight"][0].lora_A.shape} + Lora_B: {layer.parametrizations["weight"][0].lora_B.shape}'
    )
# 非LoRA パラメータのカウントは元のネットワークと一致する必要があります
assert total_parameters_non_lora == total_parameters_original
print(f'元のパラメータの総数: {total_parameters_non_lora:,}')
print(f'元のパラメータ + LoRA の総数: {total_parameters_lora + total_parameters_non_lora:,}')
print(f'LoRA によって導入されたパラメータ: {total_parameters_lora:,}')
parameters_incremment = (total_parameters_lora / total_parameters_non_lora) * 100
print(f'パラメータの増加: {parameters_incremment:.3f}%')

# 非LoRA パラメータを凍結します
for name, param in net.named_parameters():
    if 'lora' not in name:
        print(f'非LoRA パラメータ {name} を凍結します')
        param.requires_grad = False

# 数字 9 のサンプルのみを保持します
mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
digit_9_indices = mnist_trainset.targets == 9
mnist_trainset.data = mnist_trainset.data[digit_9_indices]
mnist_trainset.targets = mnist_trainset.targets[digit_9_indices]

# データローダーを作成します
train_loader = torch.utils.data.DataLoader(mnist_trainset, batch_size=10, shuffle=True)

# モデルを微調整します。100 バッチのみをトレーニングします
train(train_loader, net, epochs=1, total_iterations_limit=100)

assert torch.all(net.linear1.parametrizations.weight.original == original_weights['linear1.weight'])
assert torch.all(net.linear2.parametrizations.weight.original == original_weights['linear2.weight'])
assert torch.all(net.linear3.parametrizations.weight.original == original_weights['linear3.weight'])

enable_disable_lora(enabled=True)
# 新しい linear1.weight は LoRA パラメータ化の「forward」関数によって得られます
# 元の重みは net.linear1.parametrizations.weight.original に移動されています
# 詳細はこちら： https://pytorch.org/tutorials/intermediate/parametrizations.html#inspecting-a-parametrized-module
assert torch.equal(net.linear1.weight, net.linear1.parametrizations.weight.original + (net.linear1.parametrizations.weight[0].lora_B @ net.linear1.parametrizations.weight[0].lora_A) * net.linear1.parametrizations.weight[0].scale)

enable_disable_lora(enabled=False)
# LoRA を無効にすると、linear1.weight は元のものになります
assert torch.equal(net.linear1.weight, original_weights['linear1.weight'])

# LoRA を有効にしてテストします
enable_disable_lora(enabled=True)
test()