A head is a regular `torch.nn.Module` that can be attached to a backbone.

class ImageClassificationHead[source]

ImageClassificationHead() :: BasicModule

Abstract class for ImageClassification Heads

ImageClassificationHead.hypers[source]

Returns list of parameters like lr and wd for each param group

FullyConnectedHead

class FullyConnectedHead[source]

FullyConnectedHead(input_shape:ShapeSpec, num_classes:int, pool_type:str='avg', drop_rate:float=0.0, use_conv:bool=False, lr:float=0.002, wd:float=0, filter_wd:bool=False) :: ImageClassificationHead

Classifier head w/ configurable global pooling and dropout. From - https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/classifier.py

Arguments to FullyConnectedHead:

  • input_shape (ShapeSpec): input shape
  • num_classes (int): Number of classes for the head.
  • pool_type (str): The pooling layer to use. Check here.
  • drop_rate (float): If >0.0 then applies dropout between the pool_layer and the fc layer.
  • use_conv (bool): Use a convolutional layer as the final fc layer.
  • lr (float): Learning rate for the modules.
  • wd (float): Weight decay for the modules.
  • filter_wd (bool): Filter out bias, bn from weight_decay.
input_shape = ShapeSpec(channels=512)
tst = FullyConnectedHead(input_shape, 10)
tst
FullyConnectedHead(
  (global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=True)
  (fc): Linear(in_features=512, out_features=10, bias=True)
)

Dataclass

class FCHeadDataClass[source]

FCHeadDataClass(num_classes:int='???', pool_type:str='avg', drop_rate:float=0.0, use_conv:bool=False, lr:float=0.001, wd:float=0.0)

Base config for FullyConnectedHead

input_shape = ShapeSpec(channels=512)
c = OmegaConf.structured(FCHeadDataClass(num_classes=10))
tst = FullyConnectedHead.from_config_dict(c, input_shape=input_shape)
tst
FullyConnectedHead(
  (global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=True)
  (fc): Linear(in_features=512, out_features=10, bias=True)
)

FastaiHead

class FastaiHead[source]

FastaiHead(input_shape:ShapeSpec, num_classes:int, act:str='ReLU', lin_ftrs:Optional[List]=None, ps:Union[List, int]=0.5, concat_pool:bool=True, first_bn:bool=True, bn_final:bool=False, lr:float=0.002, wd:float=0, filter_wd:bool=False) :: ImageClassificationHead

Model head that takes in_planes features, runs through lin_ftrs, and out num_classes classes.

From - https://github.com/fastai/fastai/blob/8b1da8765fc07f1232c20fa8dc5e909d2835640c/fastai/vision/learner.py#L76

The head begins with AdaptiveConcatPool2d if concat_pool=True otherwise, it uses traditional average pooling. Then it uses a Flatten layer before going on blocks of BatchNorm, Dropout and Linear layers.

Those blocks start at in_planes, then every element of lin_ftrs (defaults to [512]) and end at num_classes. ps is a list of probabilities used for the dropouts (if you only pass 1, it will use half the value then that value as many times as necessary).

Arguments to FastaiHead:

  • input_shape (ShapeSpec): input shape
  • num_classes (int): Number of classes for the head.
  • act (str): name of the activation function to use. If None uses the default activations else the name must be in ACTIVATION_REGISTRY. Activation layers are used after every block (BatchNorm, Dropout and Linear layers) if it is not the last block.
  • lin_ftrs (List): Features of the Linear layers. (defaults to [512])
  • ps (List): list of probabilities used for the dropouts.
  • concat_pool (bool): Wether to use AdaptiveConcatPool2d or AdaptiveAveragePool2d.
  • first_bn (bool): BatchNorm Layer after pool.
  • bn_final (bool): Final Layer is BatchNorm.
  • lr (float): Learning rate for the modules.
  • wd (float): Weight decay for the modules.
  • filter_wd (bool): Filter out bias, bn from weight_decay.
input_shape = ShapeSpec(channels=512)
tst = FastaiHead(input_shape=input_shape, num_classes=10)
tst
FastaiHead(
  (layers): Sequential(
    (0): SelectAdaptivePool2d (pool_type=catavgmax, flatten=True)
    (1): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): Dropout(p=0.25, inplace=False)
    (3): Linear(in_features=1024, out_features=512, bias=False)
    (4): ReLU(inplace=True)
    (5): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (6): Dropout(p=0.5, inplace=False)
    (7): Linear(in_features=512, out_features=10, bias=False)
  )
)

Dataclass

class FastaiHeadDataClass[source]

FastaiHeadDataClass(num_classes:int='???', act:str='ReLU', lin_ftrs:Optional[List]=None, ps:Any=0.5, concat_pool:bool=True, first_bn:bool=True, bn_final:bool=False, lr:float=0.002, wd:float=0, filter_wd:bool=False)

FastaiHeadDataClass(num_classes:int='???', act:str='ReLU', lin_ftrs:Union[List, NoneType]=None, ps:Any=0.5, concat_pool:bool=True, first_bn:bool=True, bn_final:bool=False, lr:float=0.002, wd:float=0, filter_wd:bool=False)

conf = OmegaConf.structured(FastaiHeadDataClass(num_classes=10))
tst = FastaiHead.from_config_dict(conf, input_shape=input_shape)
tst
FastaiHead(
  (layers): Sequential(
    (0): SelectAdaptivePool2d (pool_type=catavgmax, flatten=True)
    (1): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): Dropout(p=0.25, inplace=False)
    (3): Linear(in_features=1024, out_features=512, bias=False)
    (4): ReLU(inplace=True)
    (5): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (6): Dropout(p=0.5, inplace=False)
    (7): Linear(in_features=512, out_features=10, bias=False)
  )
)