Source code for holocron.nn.modules.dropblock

# Copyright (C) 2019-2024, François-Guillaume Fernandez.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0> for full license details.

import torch.nn as nn
from torch import Tensor

from .. import functional as F

__all__ = ["DropBlock2d"]



[docs]
class DropBlock2d(nn.Module):
    """Implements the DropBlock module from `"DropBlock: A regularization method for convolutional networks"
    <https://arxiv.org/pdf/1810.12890.pdf>`_

    .. image:: https://github.com/frgfm/Holocron/releases/download/v0.1.3/dropblock.png
        :align: center

    Args:
        p (float, optional): probability of dropping activation value
        block_size (int, optional): size of each block that is expended from the sampled mask
        inplace (bool, optional): whether the operation should be done inplace
    """

    def __init__(self, p: float = 0.1, block_size: int = 7, inplace: bool = False) -> None:
        super().__init__()
        self.p = p
        self.block_size = block_size
        self.inplace = inplace

    @property
    def drop_prob(self) -> float:
        return self.p / self.block_size**2

    def forward(self, x: Tensor) -> Tensor:
        return F.dropblock2d(x, self.drop_prob, self.block_size, self.inplace, self.training)

    def extra_repr(self) -> str:
        return f"p={self.p}, block_size={self.block_size}, inplace={self.inplace}"