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: int, in_plane: int)

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: int, in_plane: int)

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.

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: List[int], activation_func: str, use_dropout: Optional[bool] = None)

This head defines a 3d average pooling layer (torch.nn.AvgPool3d) 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: List[int], activation_func: str, use_dropout: Optional[bool] = None)

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 – Kernel size for the 3d pooling layer.

  • 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.

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.

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)

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().