zamba.models.efficientnet_models

Classes

TimeDistributedEfficientNet

Bases: ZambaVideoClassificationLightningModule

Source code in /home/runner/work/zamba/zamba/zamba/models/efficientnet_models.py

@register_model
class TimeDistributedEfficientNet(ZambaVideoClassificationLightningModule):
    _default_model_name = (
        "time_distributed"  # used to look up default configuration for checkpoints
    )

    def __init__(
        self,
        num_frames=16,
        finetune_from: Optional[Union[os.PathLike, str]] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if finetune_from is None:
            efficientnet = timm.create_model("efficientnetv2_rw_m", pretrained=True)
            efficientnet.classifier = nn.Identity()
        else:
            efficientnet = self.from_disk(finetune_from).base.module

        # freeze base layers
        for param in efficientnet.parameters():
            param.requires_grad = False

        num_backbone_final_features = efficientnet.num_features

        self.backbone = torch.nn.ModuleList(
            [
                efficientnet.get_submodule("blocks.5"),
                efficientnet.conv_head,
                efficientnet.bn2,
                efficientnet.global_pool,
            ]
        )

        self.base = TimeDistributed(efficientnet, tdim=1)
        self.classifier = nn.Sequential(
            nn.Linear(num_backbone_final_features, 256),
            nn.Dropout(0.2),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.Flatten(),
            nn.Linear(64 * num_frames, self.num_classes),
        )

        self.save_hyperparameters("num_frames")

    def forward(self, x):
        self.base.eval()
        x = self.base(x)
        return self.classifier(x)

Attributes

backbone = torch.nn.ModuleList([efficientnet.get_submodule('blocks.5'), efficientnet.conv_head, efficientnet.bn2, efficientnet.global_pool]) instance-attribute

base = TimeDistributed(efficientnet, tdim=1) instance-attribute

classifier = nn.Sequential(nn.Linear(num_backbone_final_features, 256), nn.Dropout(0.2), nn.ReLU(), nn.Linear(256, 64), nn.Flatten(), nn.Linear(64 * num_frames, self.num_classes)) instance-attribute

lr = lr instance-attribute

model_class = type(self).__name__ instance-attribute

num_classes = len(species) instance-attribute

scheduler = torch.optim.lr_scheduler.__dict__[scheduler] instance-attribute

scheduler_params = scheduler_params instance-attribute

species = species instance-attribute

test_step_outputs = [] instance-attribute

training_step_outputs = [] instance-attribute

validation_step_outputs = [] instance-attribute

Functions

__init__(num_frames = 16, finetune_from: Optional[Union[os.PathLike, str]] = None, **kwargs: Optional[Union[os.PathLike, str]])

Source code in /home/runner/work/zamba/zamba/zamba/models/efficientnet_models.py

def __init__(
    self,
    num_frames=16,
    finetune_from: Optional[Union[os.PathLike, str]] = None,
    **kwargs,
):
    super().__init__(**kwargs)

    if finetune_from is None:
        efficientnet = timm.create_model("efficientnetv2_rw_m", pretrained=True)
        efficientnet.classifier = nn.Identity()
    else:
        efficientnet = self.from_disk(finetune_from).base.module

    # freeze base layers
    for param in efficientnet.parameters():
        param.requires_grad = False

    num_backbone_final_features = efficientnet.num_features

    self.backbone = torch.nn.ModuleList(
        [
            efficientnet.get_submodule("blocks.5"),
            efficientnet.conv_head,
            efficientnet.bn2,
            efficientnet.global_pool,
        ]
    )

    self.base = TimeDistributed(efficientnet, tdim=1)
    self.classifier = nn.Sequential(
        nn.Linear(num_backbone_final_features, 256),
        nn.Dropout(0.2),
        nn.ReLU(),
        nn.Linear(256, 64),
        nn.Flatten(),
        nn.Linear(64 * num_frames, self.num_classes),
    )

    self.save_hyperparameters("num_frames")

aggregate_step_outputs(outputs: Dict[str, np.ndarray]) -> Tuple[np.ndarray, np.ndarray, np.ndarray] staticmethod

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

@staticmethod
def aggregate_step_outputs(
    outputs: Dict[str, np.ndarray]
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    y_true = np.vstack([output["y_true"] for output in outputs])
    y_pred = np.vstack([output["y_pred"] for output in outputs])
    y_proba = np.vstack([output["y_proba"] for output in outputs])

    return y_true, y_pred, y_proba

compute_and_log_metrics(y_true: np.ndarray, y_pred: np.ndarray, y_proba: np.ndarray, subset: str)

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def compute_and_log_metrics(
    self, y_true: np.ndarray, y_pred: np.ndarray, y_proba: np.ndarray, subset: str
):
    self.log(
        f"{subset}_macro_f1",
        f1_score(y_true, y_pred, average="macro", zero_division=0),
    )

    # if only two classes, skip top_k accuracy since not enough classes
    if self.num_classes > 2:
        for k in DEFAULT_TOP_K:
            if k < self.num_classes:
                self.log(
                    f"{subset}_top_{k}_accuracy",
                    top_k_accuracy_score(
                        y_true.argmax(
                            axis=1
                        ),  # top k accuracy only supports single label case
                        y_proba,
                        labels=np.arange(y_proba.shape[1]),
                        k=k,
                    ),
                )
    else:
        self.log(f"{subset}_accuracy", accuracy_score(y_true, y_pred))

    for metric_name, label, metric in compute_species_specific_metrics(
        y_true, y_pred, self.species
    ):
        self.log(f"species/{subset}_{metric_name}/{label}", metric)

configure_optimizers()

Setup the Adam optimizer. Note, that this function also can return a lr scheduler, which is usually useful for training video models.

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def configure_optimizers(self):
    """
    Setup the Adam optimizer. Note, that this function also can return a lr scheduler, which is
    usually useful for training video models.
    """
    optim = self._get_optimizer()

    if self.scheduler is None:
        return optim
    else:
        return {
            "optimizer": optim,
            "lr_scheduler": self.scheduler(
                optim, **({} if self.scheduler_params is None else self.scheduler_params)
            ),
        }

forward(x)

Source code in /home/runner/work/zamba/zamba/zamba/models/efficientnet_models.py

def forward(self, x):
    self.base.eval()
    x = self.base(x)
    return self.classifier(x)

from_disk(path: os.PathLike, **kwargs: os.PathLike) classmethod

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

@classmethod
def from_disk(cls, path: os.PathLike, **kwargs):
    # note: we always load models onto CPU; moving to GPU is handled by `devices` in pl.Trainer
    return cls.load_from_checkpoint(path, map_location="cpu", **kwargs)

on_test_epoch_end()

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def on_test_epoch_end(self):
    y_true, y_pred, y_proba = self.aggregate_step_outputs(self.test_step_outputs)
    self.compute_and_log_metrics(y_true, y_pred, y_proba, subset="test")
    self.test_step_outputs.clear()  # free memory

on_train_start()

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def on_train_start(self):
    metrics = {"val_macro_f1": {}}

    if self.num_classes > 2:
        metrics.update(
            {f"val_top_{k}_accuracy": {} for k in DEFAULT_TOP_K if k < self.num_classes}
        )
    else:
        metrics.update({"val_accuracy": {}})

    # write hparams to hparams.yaml file, log metrics to tb hparams tab
    self.logger.log_hyperparams(self.hparams, metrics)

on_validation_epoch_end()

Aggregates validation_step outputs to compute and log the validation macro F1 and top K metrics.

Parameters:

Name	Type	Description	Default
outputs	List[dict]	list of output dictionaries from each validation step containing y_pred and y_true.	required

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def on_validation_epoch_end(self):
    """Aggregates validation_step outputs to compute and log the validation macro F1 and top K
    metrics.

    Args:
        outputs (List[dict]): list of output dictionaries from each validation step
            containing y_pred and y_true.
    """
    y_true, y_pred, y_proba = self.aggregate_step_outputs(self.validation_step_outputs)
    self.compute_and_log_metrics(y_true, y_pred, y_proba, subset="val")
    self.validation_step_outputs.clear()  # free memory

predict_step(batch, batch_idx, dataloader_idx: Optional[int] = None)

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def predict_step(self, batch, batch_idx, dataloader_idx: Optional[int] = None):
    x, y = batch
    y_hat = self(x)
    pred = torch.sigmoid(y_hat).cpu().numpy()
    return pred

test_step(batch, batch_idx)

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def test_step(self, batch, batch_idx):
    output = self._val_step(batch, batch_idx)
    self.test_step_outputs.append(output)
    return output

to_disk(path: os.PathLike)

Save out model weights to a checkpoint file on disk.

Note: this does not include callbacks, optimizer_states, or lr_schedulers. To include those, use Trainer.save_checkpoint() instead.

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def to_disk(self, path: os.PathLike):
    """Save out model weights to a checkpoint file on disk.

    Note: this does not include callbacks, optimizer_states, or lr_schedulers.
    To include those, use `Trainer.save_checkpoint()` instead.
    """

    checkpoint = {
        "state_dict": self.state_dict(),
        "hyper_parameters": self.hparams,
        "global_step": self.global_step,
        "pytorch-lightning_version": pl.__version__,
    }
    torch.save(checkpoint, path)

training_step(batch, batch_idx)

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self(x)
    loss = F.binary_cross_entropy_with_logits(y_hat, y)
    self.log("train_loss", loss.detach())
    self.training_step_outputs.append(loss)
    return loss

validation_step(batch, batch_idx)

Source code in /home/runner/work/zamba/zamba/zamba/pytorch_lightning/utils.py

def validation_step(self, batch, batch_idx):
    output = self._val_step(batch, batch_idx)
    self.validation_step_outputs.append(output)
    return output

Functions