Segmentation

class torchok.tasks.segmentation.SegmentationTask(hparams: DictConfig, backbone_name: str, head_name: str, neck_name: str, backbone_params: Optional[dict] = None, neck_params: Optional[dict] = None, head_params: Optional[dict] = None, **kwargs)

Bases: BaseTask

A class for image segmentation task.

__init__(hparams: DictConfig, backbone_name: str, head_name: str, neck_name: str, backbone_params: Optional[dict] = None, neck_params: Optional[dict] = None, head_params: Optional[dict] = None, **kwargs)

Init SegmentationTask.

Parameters

  • hparams – Hyperparameters that set in yaml file.

  • backbone_name – name of the backbone architecture in the BACKBONES registry.

  • neck_name – name of the head architecture in the DETECTION_NECKS registry.

  • head_name – name of the neck architecture in the HEADS registry.

  • backbone_params – parameters for backbone constructor.

  • neck_params – parameters for neck constructor. in_channels will be set automatically based on backbone.

  • head_params – parameters for head constructor. in_channels will be set automatically based on neck.

  • inputs – information about input model shapes and dtypes.

forward(x: Tensor) → Tensor

Forward method.

Parameters

x – torch.Tensor of shape [B, C, H, W]. Batch of input images

Returns

torch.Tensor of shape [B, num_classes, H, W], representing logits masks per each image.

forward_with_gt(batch: Dict[str, Union[Tensor, int]]) → Dict[str, Tensor]

Forward with ground truth labels.

Parameters

batch –

Dictionary with the following keys and values:

  • image (torch.Tensor):

    tensor of shape (B, C, H, W), representing input images.

  • target (torch.Tensor):

    tensor of shape [B, H, W], target class or labels masks per each image.

Returns

Dictionary with the following keys and values

  • ’prediction’: torch.Tensor of shape [B, num_classes], representing logits masks per each image.

  • ’target’: torch.Tensor of shape [B, H, W], target class or labels masks per each image. May absent.

as_module() → Sequential

Method for model representation as sequential of modules (required for checkpointing).

training: bool
precision: Union[int, str]
prepare_data_per_node: bool
allow_zero_length_dataloader_with_multiple_devices: bool