The simplest is to just use:
class SingleModel(BasicModel):
def get_layer_groups(self): return [self.model]
nlp.py has a good example (TextModel) if you want to break it up into different layers:
class TextModel(BasicModel):
def get_layer_groups(self):
m = self.model[0]
return [m.encoder, *zip(m.rnns, m.dropouths), (self.model[1], m.dropouti)]
I personally find it a little confusing to have two classes named model so I call it a ModelWrapper. Here’s another example from a method I created.
class DAEModelWrapper(BasicModel):
def init(self, model, name=‘DAE’):
self.model,self.name = model,name
def get_layer_groups(self):
m = self.model
return [m.base.embs, [m.base.bn]+children(m.base.lins)+children(m.base.bns),
children(m.head.lins)+children(m.head.bns)]
You might want to check the other forum though as there’s a thread on interacting with the different layers there.
Hope this helped.