# How can I only send trainable parameters to the client?

**URL:** <https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161>\
**Category:** Flower Help - Beginners\
**Tags:** pytorch\
**Created:** [March 20, 2024, 8:20am UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161 "2024-03-20T08:20:40Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![devmonkeyy](https://avatars.discourse-cdn.com/v4/letter/d/3ec8ea/32.png) [@devmonkeyy](https://discuss.flower.ai/u/devmonkeyy)\
**Post date:** [March 20, 2024, 8:20am UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161/1 "2024-03-20T08:20:40Z")

</div>

Hi all, I’m currently running prompt tuning on a a Flower system and having trouble finding a way to send only the trainable parameters(8 virtual tokens) to be trained in each client.

Here is my current Client class:  
class Client(fl.client.NumPyClient):  
def get\_parameters(self, config):  
return [val.cpu().numpy() for \_, val in model.state\_dict().items()]

```
    def set_parameters(self, parameters):
        params_dict = zip(model.state_dict().keys(), parameters)
        state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
        model.load_state_dict(state_dict, strict=True)

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        print("Training Started...")
        train_loop(model,
                   train_dataloader,
                   optimizer,
                   lr_scheduler,
                   DEVICE)
        print("Trianing Finished.")
        return self.get_parameters(config={}), len(train_dataloader), {}
    
    def evaluate(self, parameters, config):
        self.set_parameters(parameters)
        loss, accuracy = test_loop(model,
                                   valid_dataloader,
                                   DEVICE)
        return float(loss), len(valid_dataloader), {"accuracy": float(accuracy)}

fl.client.start_numpy_client(server_address="127.0.0.1:8080", client=Client())

```

Thanks!

---

<div class="post-metadata">

**Author:** ![javier](https://dub1.discourse-cdn.com/flex013/user_avatar/discuss.flower.ai/javier/32/116_2.png) [@javier](https://discuss.flower.ai/u/javier)\
**Post date:** [March 20, 2024, 8:28pm UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161/2 "2024-03-20T20:28:03Z")

</div>

Hey @devmonkeyy, interesting question!

As of flower 1.7 (the current version) what I’d recommend doing is to change the `get_parameters()` method so it returns just the _trainable_ parameters in the model. As you are using PyTorch I assume that _trainable_ means that _requires\_grad_=True.

The adjusted method would look as follows:

```auto
def get_parameters(model):
    return [p.cpu().detach().numpy() for p in model.parameters() if p.requires_grad]

```

You might need to adjust the code a bit to make it work with your client definition. But below I show a simple example where I take a standard ResNet18 model and then i freeze everything except the output fully connected layer.

```auto
import torch
from torchvision.models import resnet18

model = resnet18()
# Freeze everything
model = model.requires_grad_(False)
# Unfreeze output head (so it can be finetuned)
model.fc = model.fc.requires_grad_(True)

# Let's check that indeed on the the FC layer are _trainable_
for name, p in model.named_parameters():
    if p.requires_grad:
        print(f"{name} ---> shape: {p.shape}")

# prints
# fc.weight ---> shape: torch.Size([1000, 512])
# fc.bias ---> shape: torch.Size([1000])

# Let's now call the new `get_parameters` function
ndarrays = get_parameters(model)
for ndarray in ndarrays:
    print(ndarray.shape)
# prints
# (1000, 512)
# (1000,)

```

As you can see with the new function, only the _trainable_ parameters will be extracted from the model.

---

<div class="post-metadata">

**Author:** ![devmonkeyy](https://avatars.discourse-cdn.com/v4/letter/d/3ec8ea/32.png) [@devmonkeyy](https://discuss.flower.ai/u/devmonkeyy)\
**Post date:** [March 21, 2024, 3:50am UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161/3 "2024-03-21T03:50:17Z")

</div>

Thanks for the response!

I changed the get\_parameters function and froze layers in the model. When I run the code I receive the error:

> size mismatch for base\_model.bert.embeddings.word\_embeddings.weight: copying a param with shape torch.Size([8, 768]) from checkpoint, the shape in current model is torch.Size([30522, 768])

This is from the line in setParameters()

> model.load\_state\_dict(state\_dict, strict=True)

Does this mean that I have to add the frozen layers back when I set the parameters of the model?

---

<div class="post-metadata">

**Author:** ![javier](https://dub1.discourse-cdn.com/flex013/user_avatar/discuss.flower.ai/javier/32/116_2.png) [@javier](https://discuss.flower.ai/u/javier)\
**Post date:** [March 21, 2024, 9:44am UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161/4 "2024-03-21T09:44:02Z")

</div>

If you are not modifying the rest of the model during training (e.g. because maybe you initialized it with some pre-trained weights), you could then update the `set_parameters()` function to only load the state\_dict of those layers whose weights are updated.

For example, in the [ViT-finetuning](https://github.com/adap/flower/tree/main/examples/vit-finetune) example, we only finetune the last layers of the model, while keeping the rest intact. You can see [here](https://github.com/adap/flower/blob/4041299b4d75eb478a65e71ac21fe7e3f77ec2fb/examples/vit-finetune/model.py#L28) how we updated the `get_parameters()`.

```auto
def set_parameters(model, parameters):
    # `parameters` contain updated weights for the model heads
    # so let's just load those into the state_dict.
    finetune_layers = model.heads
    params_dict = zip(finetune_layers.state_dict().keys(), parameters)
    state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
    finetune_layers.load_state_dict(state_dict, strict=True)

```

So, I’d recommend doing something like this. You’ll need to find which are the layers that are _trainable_ and then construct the _state\_dict_. Does this work for you?

---

<div class="post-metadata">

**Author:** ![devmonkeyy](https://avatars.discourse-cdn.com/v4/letter/d/3ec8ea/32.png) [@devmonkeyy](https://discuss.flower.ai/u/devmonkeyy)\
**Post date:** [March 21, 2024, 8:57pm UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161/5 "2024-03-21T20:57:14Z")

</div>

Thanks for the explanation. I was able to get the code to work by only loading the state dict of the model’s prompt\_encoder. Thanks a lot for the help!

---

<div class="post-metadata">

**Author:** ![system](https://europe1.discourse-cdn.com/flex013/uploads/flower/original/1X/bc7fe36843c391cbd0b4e8b6f48275961d356abc.png) [@system](https://discuss.flower.ai/u/system)\
**Post date:** [March 28, 2024, 8:57pm UTC](https://discuss.flower.ai/t/how-can-i-only-send-trainable-parameters-to-the-client/161/6 "2024-03-28T20:57:16Z")

</div>

This topic was automatically closed 7 days after the last reply. New replies are no longer allowed.