Since both the two attr 'encoder' and 'head' are required for the model used by MOON, and few existing benchmarks support this kind of settings, here I discuss how to customize models for MOON. One can specify the model by the keyword 'model' in flgo.init(...), and the passed model should be converted to the format that is suitable for FLGo by the API flgo.conver_model(model_class:nn.Module, model_name:str). An example is as follows. Other methods that have specific requirements on the model architectures or reguarization can be implemented in this way.
This function is updated in flgo-v0.1.4, please update the library by 'pip install --upgrade flgo' into the latest version before running the example.
Remark: the class Model should define the encoder and head.
import flgo
import flgo.benchmark.partition as fbp
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, padding=2),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Flatten(1),
nn.Linear(3136, 512),
nn.ReLU(),
)
self.head = nn.Linear(512, 10)
def forward(self, x):
x = self.encoder(x)
x = self.head(x)
return x
model = flgo.convert_model(Model, 'mycnn')
moon = flgo.download_resource('.', 'moon', 'algorithm')
mnist = flgo.download_resource('.', 'mnist_classification', 'benchmark')
task = './my_task'
flgo.gen_task_by_(mnist, fbp.IIDPartitioner(num_clients=100), task)
runner = flgo.init(task, moon, option={"gpu":0, 'local_test':True, }, model=model)
runner.run()
Since both the two attr 'encoder' and 'head' are required for the model used by MOON, and few existing benchmarks support this kind of settings, here I discuss how to customize models for MOON. One can specify the model by the keyword 'model' in flgo.init(...), and the passed model should be converted to the format that is suitable for FLGo by the API flgo.conver_model(model_class:nn.Module, model_name:str). An example is as follows. Other methods that have specific requirements on the model architectures or reguarization can be implemented in this way.
This function is updated in flgo-v0.1.4, please update the library by 'pip install --upgrade flgo' into the latest version before running the example.
Remark: the class Model should define the encoder and head.