cyto_dl.nn.head.vic_reg module

class cyto_dl.nn.head.vic_reg.VICRegHead(loss, dims: List[int] = [2048, 8192], hidden_layers: List[int] = [8192, 8192])

Bases: BaseHead

Parameters:

• loss – Loss function for task

• dims – input and output dimensions for Projectornetwork

• hidden_layers – hidden layers for MLP

run_head(backbone_features, batch, stage, n_postprocess=1, run_forward=True, y_hat=None)

Run head on backbone features, calculate loss, postprocess and save image, and calculate metrics.