neuronets / nobrainer_training_scripts

1 stars 2 forks source link

Write train decorator for models #16

Closed hvgazula closed 4 months ago

hvgazula commented 4 months ago

looking at brainy_trainj.py and kwyk_train.py, pretty much the entire code as it pertains to training is similar except for the part where the model is initialized. So, it makes sense to pull apart the common code into a separate function/wrapper and simply add that as a decorator. For example

before_code
model = Segmentation(unet, ...)
model.fit()
after_code

before_code
model = Segmentation(kwyk, ...)
model.fit()
after_code

can be written as

def my_decorator(func):
    def train_wrapper(*args, **kwargs):
        before_code
        out = func(*args, **kwargs)
        after_code
        return out
    return train_wrapper

@my_decorator
def my_model():
    model = Segmentation(unet, ...)

I maybe missing a line or two but that will be the skeleton, this way, all new models can be tested with the variable components moved to config.yml [currently all models are trained on kwyk dataset, but as you can see, this choice can be moved into config.yml]