Heads

class classy_vision.heads.ClassyHead(unique_id: Optional[str] = None, num_classes: Optional[int] = None)

Base class for heads that can be attached to ClassyModel.

A head is a regular torch.nn.Module that can be attached to a pretrained model. This enables a form of transfer learning: utilizing a model trained for one dataset to extract features that can be used for other problems. A head must be attached to a models.ClassyBlock within a models.ClassyModel.

__init__(unique_id: Optional[str] = None, num_classes: Optional[int] = None)

Constructs a ClassyHead.

Parameters

  • unique_id – A unique identifier for the head. Multiple instances of the same head might be attached to a model, and unique_id is used to refer to them.

  • num_classes – Number of classes for the head.

forward(x)

Performs inference on the head.

This is a regular PyTorch method, refer to torch.nn.Module for more details

classmethod from_config(config: Dict[str, Any]) → classy_vision.heads.classy_head.ClassyHead

Instantiates a ClassyHead from a configuration.

Parameters

config – A configuration for the ClassyHead.

Returns

A ClassyHead instance.

class classy_vision.heads.FullyConnectedHead(unique_id: str, num_classes: Optional[int], in_plane: int, conv_planes: Optional[int] = None, activation: Optional[torch.nn.modules.module.Module] = None, zero_init_bias: bool = False, normalize_inputs: Optional[str] = None)

This head defines a 2d average pooling layer (torch.nn.AdaptiveAvgPool2d) followed by a fully connected layer (torch.nn.Linear).

__init__(unique_id: str, num_classes: Optional[int], in_plane: int, conv_planes: Optional[int] = None, activation: Optional[torch.nn.modules.module.Module] = None, zero_init_bias: bool = False, normalize_inputs: Optional[str] = None)

Constructor for FullyConnectedHead

Parameters

  • unique_id – A unique identifier for the head. Multiple instances of the same head might be attached to a model, and unique_id is used to refer to them.

  • num_classes – Number of classes for the head. If None, then the fully connected layer is not applied.

  • in_plane – Input size for the fully connected layer.

  • conv_planes – If specified, applies a 1x1 convolutional layer to the input before passing it to the average pooling layer. The convolution is also followed by a BatchNorm and an activation.

  • activation – The activation to be applied after the convolutional layer. Unused if conv_planes is not specified.

  • zero_init_bias – Zero initialize the bias

  • normalize_inputs – If specified, normalize the inputs after performing average pooling using the specified method. Supports “l2” normalization.

forward(x)

Performs inference on the head.

This is a regular PyTorch method, refer to torch.nn.Module for more details

classmethod from_config(config: Dict[str, Any]) → classy_vision.heads.fully_connected_head.FullyConnectedHead

Instantiates a FullyConnectedHead from a configuration.

Parameters

config – A configuration for a FullyConnectedHead. See __init__() for parameters expected in the config.

Returns

A FullyConnectedHead instance.

class classy_vision.heads.FullyConvolutionalLinearHead(unique_id: str, num_classes: int, in_plane: int, pool_size: Optional[List[int]], activation_func: str, use_dropout: Optional[bool] = None, dropout_ratio: float = 0.5)

This head defines a 3d average pooling layer (torch.nn.AvgPool3d or torch.nn.AdaptiveAvgPool3d if pool_size is None) followed by a fully convolutional linear layer. This layer performs a fully-connected projection during training, when the input size is 1x1x1. It performs a convolutional projection during testing when the input size is larger than 1x1x1.

__init__(unique_id: str, num_classes: int, in_plane: int, pool_size: Optional[List[int]], activation_func: str, use_dropout: Optional[bool] = None, dropout_ratio: float = 0.5)

Constructor for FullyConvolutionalLinearHead.

Parameters

  • unique_id – A unique identifier for the head. Multiple instances of the same head might be attached to a model, and unique_id is used to refer to them.

  • num_classes – Number of classes for the head.

  • in_plane – Input size for the fully connected layer.

  • pool_size – Optional kernel size for the 3d pooling layer. If None, use torch.nn.AdaptiveAvgPool3d with output size (1, 1, 1).

  • activation_func – activation function to use. ‘softmax’: applies softmax on the output. ‘sigmoid’: applies sigmoid on the output.

  • use_dropout – Whether to apply dropout after the pooling layer.

  • dropout_ratio – dropout ratio.

forward(x)

Performs inference on the head.

This is a regular PyTorch method, refer to torch.nn.Module for more details

classmethod from_config(config: Dict[str, Any]) → classy_vision.heads.fully_convolutional_linear_head.FullyConvolutionalLinearHead

Instantiates a FullyConvolutionalLinearHead from a configuration.

Parameters

config – A configuration for a FullyConvolutionalLinearHead. See __init__() for parameters expected in the config.

Returns

A FullyConvolutionalLinearHead instance.

class classy_vision.heads.IdentityHead(unique_id: Optional[str] = None, num_classes: Optional[int] = None)

This head returns the input without changing it. This can be attached to a model, if the output of the model is the desired result.

forward(x)

Performs inference on the head.

This is a regular PyTorch method, refer to torch.nn.Module for more details

classmethod from_config(config: Dict[str, Any]) → classy_vision.heads.identity_head.IdentityHead

Instantiates a IdentityHead from a configuration.

Parameters

config – A configuration for a IdentityHead. See __init__() for parameters expected in the config.

Returns

A IdentityHead instance.

class classy_vision.heads.VisionTransformerHead(unique_id: str, in_plane: int, num_classes: Optional[int] = None, hidden_dim: Optional[int] = None, normalize_inputs: Optional[str] = None)
__init__(unique_id: str, in_plane: int, num_classes: Optional[int] = None, hidden_dim: Optional[int] = None, normalize_inputs: Optional[str] = None)
Parameters

  • unique_id – A unique identifier for the head

  • in_plane – Input size for the fully connected layer

  • num_classes – Number of output classes for the head

  • hidden_dim – If not None, a hidden layer with the specific dimension is added

  • normalize_inputs – If specified, normalize the inputs using the specified method. Supports “l2” normalization.

forward(x)

Performs inference on the head.

This is a regular PyTorch method, refer to torch.nn.Module for more details

classmethod from_config(config)

Instantiates a ClassyHead from a configuration.

Parameters

config – A configuration for the ClassyHead.

Returns

A ClassyHead instance.

classy_vision.heads.build_head(config)

Builds a ClassyHead from a config.

This assumes a ‘name’ key in the config which is used to determine what head class to instantiate. For instance, a config {“name”: “my_head”, “foo”: “bar”} will find a class that was registered as “my_head” (see register_head()) and call .from_config on it.

classy_vision.heads.register_head(name, bypass_checks=False)

Registers a ClassyHead subclass.

This decorator allows Classy Vision to instantiate a subclass of ClassyHead from a configuration file, even if the class itself is not part of the Classy Vision framework. To use it, apply this decorator to a ClassyHead subclass, like this:

@register_head("my_head")
class MyHead(ClassyHead):
    ...

To instantiate a head from a configuration file, see build_head().