loralib_linear_layer_wrapping_with_trainable_params.py

1import torch
2import torch.nn as nn
3import loralib as lora
4
5# 1. Define your model (Example: a simple MLP)
6class MyModel(nn.Module):
7    def __init__(self):
8        super(MyModel, self).__init__()
9        # Replace a standard nn.Linear with lora.Linear
10        # r is the rank, lora_alpha is the scaling factor
11        self.fc = lora.Linear(784, 128, r=16)
12
13model = MyModel()
14
15# 2. Mark only LoRA parameters as trainable
16lora.mark_only_lora_as_trainable(model)
17
18# 3. Training loop (standard PyTorch)
19optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
20data = torch.randn(1, 784)
21target = torch.randn(1, 128)
22
23output = model(data)
24loss = nn.MSELoss()(output, target)
25loss.backward()
26optimizer.step()
27
28# 4. Saving the model (only saves LoRA weights)
29torch.save(lora.lora_state_dict(model), "lora_weights.pt")