24 lines
458 B
Python
24 lines
458 B
Python
from federatedscope.register import register_model
|
|
|
|
|
|
# Build you torch or tf model class here
|
|
class MyNet(object):
|
|
pass
|
|
|
|
|
|
# Instantiate your model class with config and data
|
|
def ModelBuilder(model_config, local_data):
|
|
|
|
model = MyNet()
|
|
|
|
return model
|
|
|
|
|
|
def call_my_net(model_config, local_data):
|
|
if model_config.type == "mynet":
|
|
model = ModelBuilder(model_config, local_data)
|
|
return model
|
|
|
|
|
|
register_model("mynet", call_my_net)
|