Heads

terratorch.models.heads.regression_head

RegressionHead

Bases: Module

Regression head

__init__(in_channels, final_act=None, learned_upscale_layers=0, channel_list=None, batch_norm=True, dropout=0)

Constructor

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
final_act	Module | None	Final activation to be applied. Defaults to None.	None
learned_upscale_layers	int	Number of Pixelshuffle layers to create. Each upscales 2x. Defaults to 0.	0
channel_list	list[int] | None	List with number of channels for each Conv layer to be created. Defaults to None.	None
batch_norm	bool	Whether to apply batch norm. Defaults to True.	True
dropout	float	Dropout value to apply. Defaults to 0.	0

terratorch.models.heads.segmentation_head

SegmentationHead

Bases: Module

Segmentation head

__init__(in_channels, num_classes, channel_list=None, dropout=0)

Constructor

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels	required
num_classes	int	Number of output classes	required
channel_list	list[int] | None	List with number of channels for each Conv layer to be created. Defaults to None.	None
dropout	float	Dropout value to apply. Defaults to 0.	0

terratorch.models.heads.classification_head

ClassificationHead

Bases: Module

Classification head

__init__(in_dim, num_classes, dim_list=None, dropout=0, linear_after_pool=False)

Constructor

Parameters:

Name	Type	Description	Default
in_dim	int	Input dimensionality	required
num_classes	int	Number of output classes	required
dim_list	list[int] | None	List with number of dimensions for each Linear layer to be created. Defaults to None.	None
dropout	float	Dropout value to apply. Defaults to 0.	0
linear_after_pool	bool	Apply pooling first, then apply the linear layer. Defaults to False	False

---

Last update: March 24, 2025