Function set_dropout_mode

modAL/dropout.py:405–436

Function to change the mode of the dropout layers (bool: train_mode -> train or evaluation) Args: model: Pytorch model dropout_layer_indexes: Indexes of the dropout layers which should be activated Choose indices from : list(torch_model.mod

(model, dropout_layer_indexes: list, train_mode: bool)

def set_dropout_mode(model, dropout_layer_indexes: list, train_mode: bool):
    """ 
        Function to change the mode of the dropout layers (bool: train_mode -> train or evaluation)

        Args: 
            model: Pytorch model
            dropout_layer_indexes: Indexes of the dropout layers which should be activated
                Choose indices from : list(torch_model.modules())    
            train_mode: boolean, true <=> train_mode, false <=> evaluation_mode 
    """

    modules = list(model.modules())  # list of all modules in the network.

    if len(dropout_layer_indexes) != 0:
        for index in dropout_layer_indexes:
            layer = modules[index]
            if layer.__class__.__name__.startswith('Dropout'):
                if True == train_mode:
                    layer.train()
                elif False == train_mode:
                    layer.eval()
            else:
                raise KeyError(
                    "The passed index: {} is not a Dropout layer".format(index))

    else:
        for module in modules:
            if module.__class__.__name__.startswith('Dropout'):
                if True == train_mode:
                    module.train()
                elif False == train_mode:
                    module.eval()