Skip to content

Utils

Helper functions

freeze_model_by_layer_name(network, freeze_layer_name=None, freeze_bn=True, unfreeze_all=False)

Unfreeze the last layers of given network starting from the layer_name. All other layers will be frozen regardless of their initial state.

Freezing the layers means their weights become non-trainable and are not updated during training. It has purpose only when the weights are pre-trained and are not random. This trade-offs the overall performance for computational efficiency.

Parameters:

Name Type Description Default
network tf.keras.Model

network to freeze layers of

 required
freeze_layer_name str

name of the first layer to be trainable. Every layer before it will be frozen

 None
freeze_bn bool

Whether to leave all BatchNormalization layers frozen. It is observed that leaving frozen during fine-tuning gives a slightly better accuracy

 True
unfreeze_all bool

Whether to unfreeze the whole network

 False
Source code in conftrainer/training/utils.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def freeze_model_by_layer_name(network: Model, freeze_layer_name: Optional[str] = None, freeze_bn: bool = True,
                               unfreeze_all: bool = False) -> None:
    """
    Unfreeze the last layers of given network starting from the layer_name. All other layers
    will be frozen regardless of their initial state.

    Freezing the layers means their weights become non-trainable and are not updated during
    training. It has purpose only when the weights are pre-trained and are not random. This
    trade-offs the overall performance for computational efficiency.

    Parameters
    ----------
    network : tf.keras.Model
        network to freeze layers of
    freeze_layer_name : str
        name of the first layer to be trainable. Every layer before it will be frozen
    freeze_bn : bool
        Whether to leave all BatchNormalization layers frozen. It is observed that leaving
        frozen during fine-tuning gives a slightly better accuracy
    unfreeze_all : bool
        Whether to unfreeze the whole network
    """
    if unfreeze_all:
        network.trainable = True
        return
    if not freeze_layer_name:
        network.trainable = False
        return
    for layer in network.layers:
        if layer.name != freeze_layer_name:
            layer.trainable = False
        else:
            break
    if freeze_bn:
        freeze_batchnorm(network)

freeze_batchnorm(network)

Freeze all BatchNormalization layers of given network

Parameters:

Name Type Description Default
network tf.keras.Model

network to freeze

 required
Source code in conftrainer/training/utils.py 
59
60
61
62
63
64
65
66
67
68
69
70
def freeze_batchnorm(network: Model) -> None:
    """
    Freeze all BatchNormalization layers of given network

    Parameters
    ----------
    network : tf.keras.Model
        network to freeze
    """
    for layer in network.layers:
        if isinstance(layer, BatchNormalization):
            layer.trainable = False

get_loss(name='SigmoidFocalCrossEntropy', **kwargs)

Initialize a loss function

Parameters:

Name Type Description Default
name str

name of the loss function. Should be from tf.losses, keras-cv.losses or conftrainer.losses

 :SigmoidFocalCrossentropy
**kwargs dict

loss-specific keyword arguments. Ex.

 {}

loss_func = get_loss("CategoricalFocalLoss", **{"alpha" : 0.75, "gamma" : 1.25})

Returns:

Name Type Description
out tf.keras.losses.Loss

loss function
Source code in conftrainer/training/utils.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def get_loss(name: str = "SigmoidFocalCrossEntropy", **kwargs) -> tf_losses.Loss:
    """
    Initialize a loss function

    Parameters
    ----------
    name : str, default :SigmoidFocalCrossentropy
        name of the loss function. Should be from tf.losses, keras-cv.losses or conftrainer.losses
    **kwargs : dict
        loss-specific keyword arguments. Ex.

    >>> loss_func = get_loss("CategoricalFocalLoss", **{"alpha" : 0.75, "gamma" : 1.25})

    Returns
    -------
    out : tf.keras.losses.Loss
        loss function
    """
    loss_config = ObjInitConfig(name=name, args=kwargs)
    modules = [tf_losses, custom_losses]
    return create_object(config=loss_config, modules=modules)

get_lr(learning_rate, schedule=None, **kwargs)

Get learning rate to use during training

Parameters:

Name Type Description Default
learning_rate float

initial learning rate

 required
schedule Optional[str]

name of the schedule to use. If None, return learning rate as is

 None
**kwargs

schedule-specific arguments

 {}

Returns:

Name Type Description
out float, tf.keras.optimizers.Scheduler

learning rate
Source code in conftrainer/training/utils.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def get_lr(learning_rate: float, schedule: Optional[str] = None,
           **kwargs) -> Union[float, tf_schedules.LearningRateSchedule]:
    """
    Get learning rate to use during training

    Parameters
    ----------
    learning_rate: float
        initial learning rate
    schedule: str = None
        name of the schedule to use. If None, return learning rate as is
    **kwargs: dict
        schedule-specific arguments

    Returns
    -------
    out: float, tf.keras.optimizers.Scheduler
        learning rate
    """

    if schedule:
        schedule_config = ObjInitConfig(name=schedule, args={"initial_learning_rate": learning_rate,
                                                             **kwargs})
        return create_object(schedule_config, modules=[tf_schedules, custom_schedules])
    return learning_rate

get_optimizer(name, learning_rate=0.001, **kwargs)

Initialize an optimizer to use during training

Parameters:

Name Type Description Default
name str

name of the optimizer module, either from tensorflow.keras.optimizers or tensorflow_addons.optimizers

 required
learning_rate float, dict

initial learning rate to initialize the optimizer with. Either a float or a schedule

 0.001
**kwargs dict

optimizer-specific arguments. Ex.

 {}

opt = get_optimizer(name="AdamW", **{"learning_rate": 0.01, "weight_decay":0.00001})

Returns:

Name Type Description
out tf.keras.optimizers.Optimizer

an optimizer
Source code in conftrainer/training/utils.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def get_optimizer(name: str, learning_rate: float = 0.001, **kwargs) -> tf_optimizers.Optimizer:
    """
    Initialize an optimizer to use during training

    Parameters
    ----------
    name : str
        name of the optimizer module, either from tensorflow.keras.optimizers or
        tensorflow_addons.optimizers
    learning_rate : float, dict
        initial learning rate to initialize the optimizer with. Either a float or a schedule
    **kwargs : dict
        optimizer-specific arguments. Ex.

    >>> opt = get_optimizer(name="AdamW", **{"learning_rate": 0.01, "weight_decay":0.00001})

    Returns
    -------
    out : tf.keras.optimizers.Optimizer
        an optimizer
    """
    opt_config = ObjInitConfig(name=name, args={"learning_rate": learning_rate, **kwargs})
    modules = [tf_optimizers, custom_optimizers]
    return create_object(config=opt_config, modules=modules)

Main util to train a network

create_and_train(config)

Main training loop. Initialize a network and a Trainer on top of it, then train & save the network based on the given configuration

Source code in conftrainer/training/train.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def create_and_train(config: TrainConfig) -> None:
    """Main training loop. Initialize a network and a Trainer on top of it,
     then train & save the network based on the given configuration"""
    arch = getattr(models, config.net_config.framework)
    network = arch(config.net_config)
    datagens = load_all_datagens(csvs=config.data_config.csvs,
                                 shape=config.net_config.input_shape,
                                 name_col=config.data_config.name_col,
                                 clean_dataset=config.data_config.clean_dataset,
                                 trainable_classes=config.data_config.trainable_classes,
                                 data_dir=config.data_config.data_dir
                                 )
    trainer = Trainer(network=network, config=config, datagens=datagens)
    trainer.train(conf=config.pretrain_config)
    if config.finetune_config.fine_tune:
        trainer.train(conf=config.finetune_config)
    trainer.predict_and_report()

train_multibranch(config)

Initialize a MultiOutputTrainer and train a network with given configurations

Source code in conftrainer/training/train.py 
34
35
36
37
38
39
40
41
42
43
44
45
def train_multibranch(config: MultiBranchTrainConfig):
    """Initialize a MultiOutputTrainer and train a network with given configurations"""
    config.data_config.per_task_data = [SingleTaskDataConfig(name=branch.name,
                                                             classes=branch.classes,
                                                             task_type=branch.task_type)
                                        for branch in config.net_config.branches]
    datagens = create_multioutput_datagens(read_config=config.data_config, shape=config.net_config.input_shape)
    net = multibranch_network(config.net_config)
    trainer = MultiOutputTrainer(network=net, config=config, datagens=datagens)
    trainer.train(config.pretrain_config)
    if config.finetune_config.fine_tune:
        trainer.train(config.finetune_config)