Training tools

get_toy_keras_convnet(num_classes)

Basic keras convolutional classifier for toy datasets.

Parameters:

Name	Type	Description	Default
num_classes	int	Number of classes for the classification task.	required

Returns:

Type	Description
Model	tf.keras.Model: model

Source code in oodeel/utils/tf_training_tools.py

def get_toy_keras_convnet(num_classes: int) -> tf.keras.Model:
    """Basic keras convolutional classifier for toy datasets.

    Args:
        num_classes (int): Number of classes for the classification task.

    Returns:
        tf.keras.Model: model
    """
    return Sequential(
        [
            Conv2D(32, kernel_size=(3, 3), activation="relu"),
            MaxPooling2D(pool_size=(2, 2)),
            Conv2D(64, kernel_size=(3, 3), activation="relu"),
            MaxPooling2D(pool_size=(2, 2)),
            Flatten(),
            Dropout(0.5),
            Dense(num_classes, activation="softmax"),
        ]
    )

get_toy_mlp(input_shape, num_classes)

Basic keras MLP classifier for toy datasets.

Parameters:

Name	Type	Description	Default
input_shape	tuple	Input data shape.	required
num_classes	int	Number of classes for the classification task.	required

Returns:

Type	Description
Model	tf.keras.Model: model

Source code in oodeel/utils/tf_training_tools.py

def get_toy_mlp(input_shape: tuple, num_classes: int) -> tf.keras.Model:
    """Basic keras MLP classifier for toy datasets.

    Args:
        input_shape (tuple): Input data shape.
        num_classes (int): Number of classes for the classification task.

    Returns:
        tf.keras.Model: model
    """
    return tf.keras.models.Sequential(
        [
            tf.keras.layers.Input(shape=input_shape),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(300, activation="relu"),
            tf.keras.layers.Dense(150, activation="relu"),
            tf.keras.layers.Dense(num_classes, activation="softmax"),
        ]
    )

train_tf_model(train_data, model, input_shape, num_classes, batch_size=128, epochs=50, loss='sparse_categorical_crossentropy', optimizer='adam', lr_scheduler=None, learning_rate=0.001, metrics=['accuracy'], imagenet_pretrained=False, validation_data=None, save_dir=None, save_best_only=True)

Loads a model from tensorflow.python.keras.applications. If the dataset is different from imagenet, trains on provided dataset.

Parameters:

Name	Type	Description	Default
train_data	Dataset	training dataset.	required
model	Union[Model, str]	if a string is provided, must be a model from tf.keras.applications or "toy_convnet" or "toy_mlp"	required
input_shape	tuple	Shape of the input images.	required
num_classes	int	Number of output classes.	required
batch_size	int	Defaults to 128.	128
epochs	int	Defaults to 50.	50
loss	str	Defaults to "sparse_categorical_crossentropy".	'sparse_categorical_crossentropy'
optimizer	str	Defaults to "adam".	'adam'
lr_scheduler	str	("cosine" | "steps" | None). Defaults to None.	None
learning_rate	float	Defaults to 1e-3.	0.001
metrics	List[str]	Validation metrics. Defaults to ["accuracy"].	['accuracy']
imagenet_pretrained	bool	Load a model pretrained on imagenet or not. Defaults to False.	False
validation_data	Optional[Dataset]	Defaults to None.	None
save_dir	Optional[str]	Directory to save the model. Defaults to None.	None
save_best_only	bool	If False, saved model will be the last one. Defaults to True.	True

Returns:

Type	Description
Model	tf.keras.Model: Trained model

Source code in oodeel/utils/tf_training_tools.py

def train_tf_model(
    train_data: tf.data.Dataset,
    model: Union[tf.keras.Model, str],
    input_shape: tuple,
    num_classes: int,
    batch_size: int = 128,
    epochs: int = 50,
    loss: str = "sparse_categorical_crossentropy",
    optimizer: str = "adam",
    lr_scheduler: Optional[str] = None,
    learning_rate: float = 1e-3,
    metrics: List[str] = ["accuracy"],
    imagenet_pretrained: bool = False,
    validation_data: Optional[tf.data.Dataset] = None,
    save_dir: Optional[str] = None,
    save_best_only: bool = True,
) -> tf.keras.Model:
    """Loads a model from tensorflow.python.keras.applications.
    If the dataset is different from imagenet, trains on provided dataset.

    Args:
        train_data (tf.data.Dataset): training dataset.
        model (Union[tf.keras.Model, str]): if a string is provided, must be a model
            from tf.keras.applications or "toy_convnet" or "toy_mlp"
        input_shape (tuple): Shape of the input images.
        num_classes (int): Number of output classes.
        batch_size (int, optional): Defaults to 128.
        epochs (int, optional): Defaults to 50.
        loss (str, optional): Defaults to "sparse_categorical_crossentropy".
        optimizer (str, optional): Defaults to "adam".
        lr_scheduler (str, optional): ("cosine" | "steps" | None). Defaults to None.
        learning_rate (float, optional): Defaults to 1e-3.
        metrics (List[str], optional): Validation metrics. Defaults to ["accuracy"].
        imagenet_pretrained (bool, optional): Load a model pretrained on imagenet or
            not. Defaults to False.
        validation_data (Optional[tf.data.Dataset], optional): Defaults to None.
        save_dir (Optional[str], optional): Directory to save the model.
            Defaults to None.
        save_best_only (bool): If False, saved model will be the last one. Defaults to
            True.

    Returns:
        tf.keras.Model: Trained model
    """
    # get data infos from dataset
    if isinstance(train_data.element_spec, dict):
        input_id = "image"
        label_id = "label"
    else:
        input_id = 0
        label_id = -1
    if input_shape is None:
        input_shape = TFDataHandler.get_feature_shape(train_data, input_id)
    if num_classes is None:
        classes = TFDataHandler.get_feature(train_data, label_id).unique()
        num_classes = len(list(classes.as_numpy_iterator()))

    # prepare model
    if isinstance(model, tf.keras.Model):
        pass
    elif isinstance(model, str):
        if model == "toy_convnet":
            model = get_toy_keras_convnet(num_classes)
        elif model == "toy_mlp":
            model = get_toy_mlp(input_shape, num_classes)
        else:
            weights = "imagenet" if imagenet_pretrained else None
            backbone = getattr(tf.keras.applications, model)(
                include_top=False, weights=weights, input_shape=input_shape
            )

            features = tf.keras.layers.Flatten()(backbone.layers[-1].output)
            output = tf.keras.layers.Dense(
                num_classes,
                activation="softmax",
            )(features)
            model = tf.keras.Model(backbone.layers[0].input, output)

    n_samples = TFDataHandler.get_dataset_length(train_data)

    # Prepare callbacks
    model_checkpoint_callback = []

    if save_dir is not None:
        checkpoint_filepath = save_dir
        model_checkpoint_callback.append(
            tf.keras.callbacks.ModelCheckpoint(
                filepath=checkpoint_filepath,
                save_weights_only=True,
                monitor="val_accuracy",
                mode="max",
                save_best_only=save_best_only,
            )
        )

    if len(model_checkpoint_callback) == 0:
        model_checkpoint_callback = None

    # optimizer
    decay_steps = int(epochs * n_samples / batch_size)
    if lr_scheduler == "cosine":
        learning_rate_fn = tf.keras.experimental.CosineDecay(
            learning_rate, decay_steps=decay_steps
        )
    elif lr_scheduler == "steps":
        values = list(learning_rate * np.array([1, 0.1, 0.01]))
        boundaries = list(np.round(decay_steps * np.array([1 / 3, 2 / 3])).astype(int))
        learning_rate_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
            boundaries, values
        )
    else:
        learning_rate_fn = learning_rate

    config = {
        "class_name": optimizer,
        "config": {
            "learning_rate": learning_rate_fn,
        },
    }

    if optimizer == "SGD":
        config["config"]["momentum"] = 0.9
        config["config"]["decay"] = 5e-4

    keras_optimizer = tf.keras.optimizers.get(config)

    model.compile(loss=loss, optimizer=keras_optimizer, metrics=metrics)

    model.fit(
        train_data,
        validation_data=validation_data,
        epochs=epochs,
        callbacks=model_checkpoint_callback,
    )

    if save_dir is not None:
        model.load_weights(save_dir)
        model.save(save_dir)
    return model

ToyTorchConvnet

Bases: Sequential

Basic torch convolutional classifier for toy datasets.

Parameters:

Name	Type	Description	Default
input_shape	tuple	Input data shape.	required
num_classes	int	Number of classes for the classification task.	required

Source code in oodeel/utils/torch_training_tools.py

class ToyTorchConvnet(nn.Sequential):
    """Basic torch convolutional classifier for toy datasets.

    Args:
        input_shape (tuple): Input data shape.
        num_classes (int): Number of classes for the classification task.
    """

    def __init__(self, input_shape: tuple, num_classes: int):
        self.input_shape = input_shape

        # features
        features = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", nn.Conv2d(input_shape[0], 32, 3)),
                    ("relu1", nn.ReLU()),
                    ("pool1", nn.MaxPool2d(2, 2)),
                    ("conv2", nn.Conv2d(32, 64, 3)),
                    ("relu2", nn.ReLU()),
                    ("pool2", nn.MaxPool2d(2, 2)),
                    ("flatten", nn.Flatten()),
                ]
            )
        )

        # fc head
        fc_input_shape = self._calculate_fc_input_shape(features)
        fcs = nn.Sequential(
            OrderedDict(
                [
                    ("dropout", nn.Dropout(0.5)),
                    ("fc1", nn.Linear(fc_input_shape, num_classes)),
                ]
            )
        )

        # Sequential class init
        super().__init__(
            OrderedDict([*features._modules.items(), *fcs._modules.items()])
        )

    def _calculate_fc_input_shape(self, features):
        """Get tensor shape after passing a features network."""
        input_tensor = torch.ones(tuple([1] + list(self.input_shape)))
        x = features(input_tensor)
        output_size = x.view(x.size(0), -1).size(1)
        return output_size

ToyTorchMLP

Bases: Sequential

Basic torch MLP classifier for toy datasets.

Parameters:

Name	Type	Description	Default
input_shape	tuple	Input data shape.	required
num_classes	int	Number of classes for the classification task.	required

Source code in oodeel/utils/torch_training_tools.py

class ToyTorchMLP(nn.Sequential):
    """Basic torch MLP classifier for toy datasets.

    Args:
        input_shape (tuple): Input data shape.
        num_classes (int): Number of classes for the classification task.
    """

    def __init__(self, input_shape: tuple, num_classes: int):
        self.input_shape = input_shape

        # build toy mlp
        mlp_modules = OrderedDict(
            [
                ("flatten", nn.Flatten()),
                ("dense1", nn.Linear(np.prod(input_shape), 300)),
                ("relu1", nn.ReLU()),
                ("dense2", nn.Linear(300, 150)),
                ("relu2", nn.ReLU()),
                ("fc1", nn.Linear(150, num_classes)),
            ]
        )
        super().__init__(mlp_modules)

train_torch_model(train_data, model, num_classes, epochs=50, loss='CrossEntropyLoss', optimizer='Adam', lr_scheduler='cosine', learning_rate=0.001, imagenet_pretrained=False, validation_data=None, save_dir=None, cuda_idx=0)

Load a model (toy classifier or from torchvision.models) and train it over a torch dataloader.

Parameters:

Name	Type	Description	Default
train_data	DataLoader	train dataloader	required
model	Union[Module, str]	if a string is provided, must be a model from torchvision.models or "toy_convnet" or "toy_mlp.	required
num_classes	int	Number of output classes.	required
epochs	int	Defaults to 50.	50
loss	str	Defaults to "CrossEntropyLoss".	'CrossEntropyLoss'
optimizer	str	Defaults to "Adam".	'Adam'
lr_scheduler	str	("cosine" | "steps" | None). Defaults to None.	'cosine'
learning_rate	float	Defaults to 1e-3.	0.001
imagenet_pretrained	bool	Load a model pretrained on imagenet or not. Defaults to False.	False
validation_data	Optional[DataLoader]	Defaults to None.	None
save_dir	Optional[str]	Directory to save the model. Defaults to None.	None
cuda_idx	int	idx of cuda device to use. Defaults to 0.	0

Returns:

Type	Description
Module	nn.Module: trained model

Source code in oodeel/utils/torch_training_tools.py

def train_torch_model(
    train_data: DataLoader,
    model: Union[nn.Module, str],
    num_classes: int,
    epochs: int = 50,
    loss: str = "CrossEntropyLoss",
    optimizer: str = "Adam",
    lr_scheduler: str = "cosine",
    learning_rate: float = 1e-3,
    imagenet_pretrained: bool = False,
    validation_data: Optional[DataLoader] = None,
    save_dir: Optional[str] = None,
    cuda_idx: int = 0,
) -> nn.Module:
    """
    Load a model (toy classifier or from torchvision.models) and train
    it over a torch dataloader.

    Args:
        train_data (DataLoader): train dataloader
        model (Union[nn.Module, str]): if a string is provided, must be a model from
            torchvision.models or "toy_convnet" or "toy_mlp.
        num_classes (int): Number of output classes.
        epochs (int, optional): Defaults to 50.
        loss (str, optional): Defaults to "CrossEntropyLoss".
        optimizer (str, optional): Defaults to "Adam".
        lr_scheduler (str, optional): ("cosine" | "steps" | None). Defaults to None.
        learning_rate (float, optional): Defaults to 1e-3.
        imagenet_pretrained (bool, optional): Load a model pretrained on imagenet or
            not. Defaults to False.
        validation_data (Optional[DataLoader], optional): Defaults to None.
        save_dir (Optional[str], optional): Directory to save the model.
            Defaults to None.
        cuda_idx (int): idx of cuda device to use. Defaults to 0.

    Returns:
        nn.Module: trained model
    """
    # device
    device = torch.device(
        f"cuda:{cuda_idx}"
        if torch.cuda.is_available() and cuda_idx is not None
        else "cpu"
    )

    # Prepare model
    if isinstance(model, nn.Module):
        model = model.to(device)
    elif isinstance(model, str):
        if model == "toy_convnet":
            # toy model
            input_shape = tuple(next(iter(train_data))[0].shape[1:])
            model = ToyTorchConvnet(input_shape, num_classes).to(device)
        elif model == "toy_mlp":
            # toy model
            input_shape = tuple(next(iter(train_data))[0].shape[1:])
            model = ToyTorchMLP(input_shape, num_classes).to(device)
        else:
            # torchvision model
            model = getattr(torchvision.models, model)(
                num_classes=num_classes, pretrained=imagenet_pretrained
            ).to(device)

    # define optimizer and learning rate scheduler
    optimizer = getattr(optim, optimizer)(model.parameters(), lr=learning_rate)
    n_steps = len(train_data) * epochs
    if lr_scheduler == "cosine":
        lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, n_steps)
    elif lr_scheduler == "steps":
        boundaries = list(np.round(n_steps * np.array([1 / 3, 2 / 3])).astype(int))
        lr_scheduler = optim.lr_scheduler.MultiStepLR(
            optimizer, milestones=boundaries, gamma=0.1
        )

    # define loss
    criterion = getattr(nn, loss)()

    # train
    model = _train(
        model,
        train_data,
        validation_data=validation_data,
        epochs=epochs,
        criterion=criterion,
        optimizer=optimizer,
        lr_scheduler=lr_scheduler,
        save_dir=save_dir,
        device=device,
    )
    return model