mednet.models.classify.vit

Vision Transformer architecture implemented with timm.

Classes

ViT([training_step_fn, loss_type, ...])

Vision Transformer architecture implemented with timm.

class mednet.models.classify.vit.ViT(training_step_fn=default_training_step, loss_type=torch.nn.BCEWithLogitsLoss, loss_arguments=None, optimizer_type=torch.optim.Adam, optimizer_arguments=None, scheduler_type=None, scheduler_arguments=None, model_transforms=None, augmentation_transforms=None, architecture='vit_base_patch16_224.augreg_in21k', pretrained=True, img_size=224, drop_path_rate=0.0, drop_rate=0.0, global_pool='token', peft_config=None, num_classes=1)[source]

Bases: Model

Vision Transformer architecture implemented with timm.

Parameters:
  • training_step_fn (Callable) – The training step function to use.

  • loss_type (type[Module]) –

    The loss to be used for training and evaluation.

    Warning

    The loss should be set to always return batch averages (as opposed to the batch sum), as our logging system expects it so.

  • loss_arguments (dict[str, Any] | None) – Arguments to the loss.

  • optimizer_type (type[Optimizer]) – The type of optimizer to use for training.

  • optimizer_arguments (dict[str, Any] | None) – Arguments to the optimizer after params.

  • scheduler_type (type[LRScheduler] | None) – The type of scheduler to use for training.

  • scheduler_arguments (dict[str, Any] | None) – Arguments to the scheduler after params.

  • model_transforms (Optional[Sequence[Callable[[Tensor], Tensor]]]) – An optional sequence of torch modules containing transforms to be applied on the input before it is fed into the network.

  • augmentation_transforms (Optional[Sequence[Callable[[Tensor], Tensor]]]) – An optional sequence of torch modules containing transforms to be applied on the input before it is fed into the network.

  • architecture (Literal['vit_base_patch14_dinov2', 'vit_base_patch14_dinov2.lvd142m', 'vit_base_patch14_reg4_dinov2', 'vit_base_patch14_reg4_dinov2.lvd142m', 'vit_base_patch16_18x2_224', 'vit_base_patch16_18x2_224.untrained', 'vit_base_patch16_224', 'vit_base_patch16_224.augreg2_in21k_ft_in1k', 'vit_base_patch16_224.augreg_in1k', 'vit_base_patch16_224.augreg_in21k', 'vit_base_patch16_224.augreg_in21k_ft_in1k', 'vit_base_patch16_224.dino', 'vit_base_patch16_224.mae', 'vit_base_patch16_224.orig_in21k', 'vit_base_patch16_224.orig_in21k_ft_in1k', 'vit_base_patch16_224.sam_in1k', 'vit_base_patch16_224_miil', 'vit_base_patch16_224_miil.in21k', 'vit_base_patch16_224_miil.in21k_ft_in1k', 'vit_base_patch16_384', 'vit_base_patch16_384.augreg_in1k', 'vit_base_patch16_384.augreg_in21k_ft_in1k', 'vit_base_patch16_384.orig_in21k_ft_in1k', 'vit_base_patch16_clip_224', 'vit_base_patch16_clip_224.datacompxl', 'vit_base_patch16_clip_224.dfn2b', 'vit_base_patch16_clip_224.laion2b', 'vit_base_patch16_clip_224.laion2b_ft_in12k', 'vit_base_patch16_clip_224.laion2b_ft_in12k_in1k', 'vit_base_patch16_clip_224.laion2b_ft_in1k', 'vit_base_patch16_clip_224.laion400m_e32', 'vit_base_patch16_clip_224.metaclip_2pt5b', 'vit_base_patch16_clip_224.metaclip_400m', 'vit_base_patch16_clip_224.openai', 'vit_base_patch16_clip_224.openai_ft_in12k', 'vit_base_patch16_clip_224.openai_ft_in12k_in1k', 'vit_base_patch16_clip_224.openai_ft_in1k', 'vit_base_patch16_clip_384', 'vit_base_patch16_clip_384.laion2b_ft_in12k_in1k', 'vit_base_patch16_clip_384.laion2b_ft_in1k', 'vit_base_patch16_clip_384.openai_ft_in12k_in1k', 'vit_base_patch16_clip_384.openai_ft_in1k', 'vit_base_patch16_clip_quickgelu_224', 'vit_base_patch16_clip_quickgelu_224.metaclip_2pt5b', 'vit_base_patch16_clip_quickgelu_224.metaclip_400m', 'vit_base_patch16_clip_quickgelu_224.openai', 'vit_base_patch16_gap_224', 'vit_base_patch16_gap_224.untrained', 'vit_base_patch16_plus_240', 'vit_base_patch16_plus_240.untrained', 'vit_base_patch16_plus_clip_240', 'vit_base_patch16_plus_clip_240.laion400m_e32', 'vit_base_patch16_reg4_gap_256', 'vit_base_patch16_reg4_gap_256.untrained', 'vit_base_patch16_rpn_224', 'vit_base_patch16_rpn_224.sw_in1k', 'vit_base_patch16_siglip_224', 'vit_base_patch16_siglip_224.v2_webli', 'vit_base_patch16_siglip_224.webli', 'vit_base_patch16_siglip_256', 'vit_base_patch16_siglip_256.v2_webli', 'vit_base_patch16_siglip_256.webli', 'vit_base_patch16_siglip_256.webli_i18n', 'vit_base_patch16_siglip_384', 'vit_base_patch16_siglip_384.v2_webli', 'vit_base_patch16_siglip_384.webli', 'vit_base_patch16_siglip_512', 'vit_base_patch16_siglip_512.v2_webli', 'vit_base_patch16_siglip_512.webli', 'vit_base_patch16_siglip_gap_224', 'vit_base_patch16_siglip_gap_224.v2_webli', 'vit_base_patch16_siglip_gap_224.webli', 'vit_base_patch16_siglip_gap_256', 'vit_base_patch16_siglip_gap_256.v2_webli', 'vit_base_patch16_siglip_gap_256.webli', 'vit_base_patch16_siglip_gap_256.webli_i18n', 'vit_base_patch16_siglip_gap_384', 'vit_base_patch16_siglip_gap_384.v2_webli', 'vit_base_patch16_siglip_gap_384.webli', 'vit_base_patch16_siglip_gap_512', 'vit_base_patch16_siglip_gap_512.v2_webli', 'vit_base_patch16_siglip_gap_512.webli', 'vit_base_patch16_xp_224', 'vit_base_patch16_xp_224.untrained', 'vit_base_patch32_224', 'vit_base_patch32_224.augreg_in1k', 'vit_base_patch32_224.augreg_in21k', 'vit_base_patch32_224.augreg_in21k_ft_in1k', 'vit_base_patch32_224.orig_in21k', 'vit_base_patch32_224.sam_in1k', 'vit_base_patch32_384', 'vit_base_patch32_384.augreg_in1k', 'vit_base_patch32_384.augreg_in21k_ft_in1k', 'vit_base_patch32_clip_224', 'vit_base_patch32_clip_224.datacompxl', 'vit_base_patch32_clip_224.laion2b', 'vit_base_patch32_clip_224.laion2b_ft_in12k_in1k', 'vit_base_patch32_clip_224.laion2b_ft_in1k', 'vit_base_patch32_clip_224.laion400m_e32', 'vit_base_patch32_clip_224.metaclip_2pt5b', 'vit_base_patch32_clip_224.metaclip_400m', 'vit_base_patch32_clip_224.openai', 'vit_base_patch32_clip_224.openai_ft_in12k_in1k', 'vit_base_patch32_clip_224.openai_ft_in1k', 'vit_base_patch32_clip_256', 'vit_base_patch32_clip_256.datacompxl', 'vit_base_patch32_clip_384', 'vit_base_patch32_clip_384.laion2b_ft_in12k_in1k', 'vit_base_patch32_clip_384.openai_ft_in12k_in1k', 'vit_base_patch32_clip_448', 'vit_base_patch32_clip_448.laion2b_ft_in12k_in1k', 'vit_base_patch32_clip_quickgelu_224', 'vit_base_patch32_clip_quickgelu_224.laion400m_e32', 'vit_base_patch32_clip_quickgelu_224.metaclip_2pt5b', 'vit_base_patch32_clip_quickgelu_224.metaclip_400m', 'vit_base_patch32_clip_quickgelu_224.openai', 'vit_base_patch32_plus_256', 'vit_base_patch32_plus_256.untrained', 'vit_base_patch32_siglip_256', 'vit_base_patch32_siglip_256.v2_webli', 'vit_base_patch32_siglip_gap_256', 'vit_base_patch32_siglip_gap_256.v2_webli', 'vit_base_patch8_224', 'vit_base_patch8_224.augreg2_in21k_ft_in1k', 'vit_base_patch8_224.augreg_in21k', 'vit_base_patch8_224.augreg_in21k_ft_in1k', 'vit_base_patch8_224.dino', 'vit_betwixt_patch16_gap_256', 'vit_betwixt_patch16_gap_256.untrained', 'vit_betwixt_patch16_reg1_gap_256', 'vit_betwixt_patch16_reg1_gap_256.sbb_in1k', 'vit_betwixt_patch16_reg4_gap_256', 'vit_betwixt_patch16_reg4_gap_256.sbb2_e200_in12k', 'vit_betwixt_patch16_reg4_gap_256.sbb2_e200_in12k_ft_in1k', 'vit_betwixt_patch16_reg4_gap_256.sbb_in12k', 'vit_betwixt_patch16_reg4_gap_256.sbb_in12k_ft_in1k', 'vit_betwixt_patch16_reg4_gap_256.sbb_in1k', 'vit_betwixt_patch16_reg4_gap_384', 'vit_betwixt_patch16_reg4_gap_384.sbb2_e200_in12k_ft_in1k', 'vit_betwixt_patch32_clip_224', 'vit_betwixt_patch32_clip_224.tinyclip_laion400m', 'vit_dlittle_patch16_reg1_gap_256', 'vit_dlittle_patch16_reg1_gap_256.sbb_nadamuon_in1k', 'vit_dpwee_patch16_reg1_gap_256', 'vit_dpwee_patch16_reg1_gap_256.sbb_in1k', 'vit_dpwee_patch16_reg1_gap_256.sbb_nadamuon_in1k', 'vit_dwee_patch16_reg1_gap_256', 'vit_dwee_patch16_reg1_gap_256.sbb_in1k', 'vit_dwee_patch16_reg1_gap_256.sbb_nadamuon_in1k', 'vit_giant_patch14_224', 'vit_giant_patch14_224.untrained', 'vit_giant_patch14_clip_224', 'vit_giant_patch14_clip_224.laion2b', 'vit_giant_patch14_dinov2', 'vit_giant_patch14_dinov2.lvd142m', 'vit_giant_patch14_reg4_dinov2', 'vit_giant_patch14_reg4_dinov2.lvd142m', 'vit_giant_patch16_gap_224', 'vit_giant_patch16_gap_224.in22k_ijepa', 'vit_giantopt_patch16_siglip_256', 'vit_giantopt_patch16_siglip_256.v2_webli', 'vit_giantopt_patch16_siglip_384', 'vit_giantopt_patch16_siglip_384.v2_webli', 'vit_giantopt_patch16_siglip_gap_256', 'vit_giantopt_patch16_siglip_gap_256.v2_webli', 'vit_giantopt_patch16_siglip_gap_384', 'vit_giantopt_patch16_siglip_gap_384.v2_webli', 'vit_gigantic_patch14_224', 'vit_gigantic_patch14_224.untrained', 'vit_gigantic_patch14_clip_224', 'vit_gigantic_patch14_clip_224.laion2b', 'vit_gigantic_patch14_clip_224.metaclip2_worldwide', 'vit_gigantic_patch14_clip_224.metaclip_2pt5b', 'vit_gigantic_patch14_clip_378', 'vit_gigantic_patch14_clip_378.metaclip2_worldwide', 'vit_gigantic_patch14_clip_quickgelu_224', 'vit_gigantic_patch14_clip_quickgelu_224.metaclip_2pt5b', 'vit_huge_patch14_224', 'vit_huge_patch14_224.mae', 'vit_huge_patch14_224.orig_in21k', 'vit_huge_patch14_224.untrained', 'vit_huge_patch14_clip_224', 'vit_huge_patch14_clip_224.dfn5b', 'vit_huge_patch14_clip_224.laion2b', 'vit_huge_patch14_clip_224.laion2b_ft_in12k', 'vit_huge_patch14_clip_224.laion2b_ft_in12k_in1k', 'vit_huge_patch14_clip_224.laion2b_ft_in1k', 'vit_huge_patch14_clip_224.metaclip2_worldwide', 'vit_huge_patch14_clip_224.metaclip_2pt5b', 'vit_huge_patch14_clip_224.metaclip_altogether', 'vit_huge_patch14_clip_336', 'vit_huge_patch14_clip_336.laion2b_ft_in12k_in1k', 'vit_huge_patch14_clip_336.laion2b_ft_in1k', 'vit_huge_patch14_clip_378', 'vit_huge_patch14_clip_378.dfn5b', 'vit_huge_patch14_clip_378.metaclip2_worldwide', 'vit_huge_patch14_clip_quickgelu_224', 'vit_huge_patch14_clip_quickgelu_224.dfn5b', 'vit_huge_patch14_clip_quickgelu_224.metaclip2_worldwide', 'vit_huge_patch14_clip_quickgelu_224.metaclip_2pt5b', 'vit_huge_patch14_clip_quickgelu_378', 'vit_huge_patch14_clip_quickgelu_378.dfn5b', 'vit_huge_patch14_gap_224', 'vit_huge_patch14_gap_224.in1k_ijepa', 'vit_huge_patch14_gap_224.in22k_ijepa', 'vit_huge_patch14_xp_224', 'vit_huge_patch14_xp_224.untrained', 'vit_huge_patch16_gap_448', 'vit_huge_patch16_gap_448.in1k_ijepa', 'vit_intern300m_patch14_448', 'vit_intern300m_patch14_448.ogvl_2pt5', 'vit_intern300m_patch14_448.ogvl_dist', 'vit_large_patch14_224', 'vit_large_patch14_224.untrained', 'vit_large_patch14_clip_224', 'vit_large_patch14_clip_224.apple_mclip2_dfndr2b', 'vit_large_patch14_clip_224.datacompxl', 'vit_large_patch14_clip_224.dfn2b', 'vit_large_patch14_clip_224.dfn2b_s39b', 'vit_large_patch14_clip_224.laion2b', 'vit_large_patch14_clip_224.laion2b_ft_in12k', 'vit_large_patch14_clip_224.laion2b_ft_in12k_in1k', 'vit_large_patch14_clip_224.laion2b_ft_in1k', 'vit_large_patch14_clip_224.laion400m_e32', 'vit_large_patch14_clip_224.metaclip_2pt5b', 'vit_large_patch14_clip_224.metaclip_400m', 'vit_large_patch14_clip_224.openai', 'vit_large_patch14_clip_224.openai_ft_in12k', 'vit_large_patch14_clip_224.openai_ft_in12k_in1k', 'vit_large_patch14_clip_224.openai_ft_in1k', 'vit_large_patch14_clip_336', 'vit_large_patch14_clip_336.laion2b_ft_in12k_in1k', 'vit_large_patch14_clip_336.laion2b_ft_in1k', 'vit_large_patch14_clip_336.openai', 'vit_large_patch14_clip_336.openai_ft_in12k_in1k', 'vit_large_patch14_clip_quickgelu_224', 'vit_large_patch14_clip_quickgelu_224.dfn2b', 'vit_large_patch14_clip_quickgelu_224.metaclip_2pt5b', 'vit_large_patch14_clip_quickgelu_224.metaclip_400m', 'vit_large_patch14_clip_quickgelu_224.openai', 'vit_large_patch14_clip_quickgelu_336', 'vit_large_patch14_clip_quickgelu_336.openai', 'vit_large_patch14_dinov2', 'vit_large_patch14_dinov2.lvd142m', 'vit_large_patch14_reg4_dinov2', 'vit_large_patch14_reg4_dinov2.lvd142m', 'vit_large_patch14_xp_224', 'vit_large_patch14_xp_224.untrained', 'vit_large_patch16_224', 'vit_large_patch16_224.augreg_in21k', 'vit_large_patch16_224.augreg_in21k_ft_in1k', 'vit_large_patch16_224.mae', 'vit_large_patch16_224.orig_in21k', 'vit_large_patch16_384', 'vit_large_patch16_384.augreg_in21k_ft_in1k', 'vit_large_patch16_siglip_256', 'vit_large_patch16_siglip_256.v2_webli', 'vit_large_patch16_siglip_256.webli', 'vit_large_patch16_siglip_384', 'vit_large_patch16_siglip_384.v2_webli', 'vit_large_patch16_siglip_384.webli', 'vit_large_patch16_siglip_512', 'vit_large_patch16_siglip_512.v2_webli', 'vit_large_patch16_siglip_gap_256', 'vit_large_patch16_siglip_gap_256.v2_webli', 'vit_large_patch16_siglip_gap_256.webli', 'vit_large_patch16_siglip_gap_384', 'vit_large_patch16_siglip_gap_384.v2_webli', 'vit_large_patch16_siglip_gap_384.webli', 'vit_large_patch16_siglip_gap_512', 'vit_large_patch16_siglip_gap_512.v2_webli', 'vit_large_patch32_224', 'vit_large_patch32_224.orig_in21k', 'vit_large_patch32_384', 'vit_large_patch32_384.orig_in21k_ft_in1k', 'vit_little_patch16_reg1_gap_256', 'vit_little_patch16_reg1_gap_256.sbb_in12k', 'vit_little_patch16_reg1_gap_256.sbb_in12k_ft_in1k', 'vit_little_patch16_reg4_gap_256', 'vit_little_patch16_reg4_gap_256.sbb_in1k', 'vit_medium_patch16_clip_224', 'vit_medium_patch16_clip_224.tinyclip_yfcc15m', 'vit_medium_patch16_gap_240', 'vit_medium_patch16_gap_240.sw_in12k', 'vit_medium_patch16_gap_256', 'vit_medium_patch16_gap_256.sw_in12k_ft_in1k', 'vit_medium_patch16_gap_384', 'vit_medium_patch16_gap_384.sw_in12k_ft_in1k', 'vit_medium_patch16_reg1_gap_256', 'vit_medium_patch16_reg1_gap_256.sbb_in1k', 'vit_medium_patch16_reg4_gap_256', 'vit_medium_patch16_reg4_gap_256.sbb_in12k', 'vit_medium_patch16_reg4_gap_256.sbb_in12k_ft_in1k', 'vit_medium_patch16_reg4_gap_256.sbb_in1k', 'vit_medium_patch32_clip_224', 'vit_medium_patch32_clip_224.tinyclip_laion400m', 'vit_mediumd_patch16_reg4_gap_256', 'vit_mediumd_patch16_reg4_gap_256.sbb2_e200_in12k', 'vit_mediumd_patch16_reg4_gap_256.sbb2_e200_in12k_ft_in1k', 'vit_mediumd_patch16_reg4_gap_256.sbb_in12k', 'vit_mediumd_patch16_reg4_gap_256.sbb_in12k_ft_in1k', 'vit_mediumd_patch16_reg4_gap_384', 'vit_mediumd_patch16_reg4_gap_384.sbb2_e200_in12k_ft_in1k', 'vit_pwee_patch16_reg1_gap_256', 'vit_pwee_patch16_reg1_gap_256.sbb_in1k', 'vit_small_patch14_dinov2', 'vit_small_patch14_dinov2.lvd142m', 'vit_small_patch14_reg4_dinov2', 'vit_small_patch14_reg4_dinov2.lvd142m', 'vit_small_patch16_18x2_224', 'vit_small_patch16_18x2_224.untrained', 'vit_small_patch16_224', 'vit_small_patch16_224.augreg_in1k', 'vit_small_patch16_224.augreg_in21k', 'vit_small_patch16_224.augreg_in21k_ft_in1k', 'vit_small_patch16_224.dino', 'vit_small_patch16_36x1_224', 'vit_small_patch16_36x1_224.untrained', 'vit_small_patch16_384', 'vit_small_patch16_384.augreg_in1k', 'vit_small_patch16_384.augreg_in21k_ft_in1k', 'vit_small_patch32_224', 'vit_small_patch32_224.augreg_in21k', 'vit_small_patch32_224.augreg_in21k_ft_in1k', 'vit_small_patch32_384', 'vit_small_patch32_384.augreg_in21k_ft_in1k', 'vit_small_patch8_224', 'vit_small_patch8_224.dino', 'vit_so150m2_patch16_reg1_gap_256', 'vit_so150m2_patch16_reg1_gap_256.sbb_e200_in12k', 'vit_so150m2_patch16_reg1_gap_256.sbb_e200_in12k_ft_in1k', 'vit_so150m2_patch16_reg1_gap_384', 'vit_so150m2_patch16_reg1_gap_384.sbb_e200_in12k_ft_in1k', 'vit_so150m2_patch16_reg1_gap_448', 'vit_so150m2_patch16_reg1_gap_448.sbb_e200_in12k_ft_in1k', 'vit_so150m_patch16_reg4_gap_256', 'vit_so150m_patch16_reg4_gap_256.sbb_e250_in12k', 'vit_so150m_patch16_reg4_gap_256.sbb_e250_in12k_ft_in1k', 'vit_so150m_patch16_reg4_gap_384', 'vit_so150m_patch16_reg4_gap_384.sbb_e250_in12k_ft_in1k', 'vit_so150m_patch16_reg4_map_256', 'vit_so150m_patch16_reg4_map_256.untrained', 'vit_so400m_patch14_siglip_224', 'vit_so400m_patch14_siglip_224.v2_webli', 'vit_so400m_patch14_siglip_224.webli', 'vit_so400m_patch14_siglip_378', 'vit_so400m_patch14_siglip_378.v2_webli', 'vit_so400m_patch14_siglip_378.webli', 'vit_so400m_patch14_siglip_378.webli_ft_in1k', 'vit_so400m_patch14_siglip_384', 'vit_so400m_patch14_siglip_384.webli', 'vit_so400m_patch14_siglip_gap_224', 'vit_so400m_patch14_siglip_gap_224.pali2_10b_pt', 'vit_so400m_patch14_siglip_gap_224.pali2_3b_pt', 'vit_so400m_patch14_siglip_gap_224.pali_mix', 'vit_so400m_patch14_siglip_gap_224.pali_pt', 'vit_so400m_patch14_siglip_gap_224.v2_webli', 'vit_so400m_patch14_siglip_gap_224.webli', 'vit_so400m_patch14_siglip_gap_378', 'vit_so400m_patch14_siglip_gap_378.v2_webli', 'vit_so400m_patch14_siglip_gap_378.webli', 'vit_so400m_patch14_siglip_gap_378.webli_ft_in1k', 'vit_so400m_patch14_siglip_gap_384', 'vit_so400m_patch14_siglip_gap_384.webli', 'vit_so400m_patch14_siglip_gap_448', 'vit_so400m_patch14_siglip_gap_448.pali2_10b_docci', 'vit_so400m_patch14_siglip_gap_448.pali2_10b_pt', 'vit_so400m_patch14_siglip_gap_448.pali2_3b_docci', 'vit_so400m_patch14_siglip_gap_448.pali2_3b_pt', 'vit_so400m_patch14_siglip_gap_448.pali_mix', 'vit_so400m_patch14_siglip_gap_448.pali_ocrvqa', 'vit_so400m_patch14_siglip_gap_448.pali_pt', 'vit_so400m_patch14_siglip_gap_448.pali_refcoco_seg', 'vit_so400m_patch14_siglip_gap_896', 'vit_so400m_patch14_siglip_gap_896.pali2_10b_pt', 'vit_so400m_patch14_siglip_gap_896.pali2_3b_pt', 'vit_so400m_patch14_siglip_gap_896.pali_ocrvqa', 'vit_so400m_patch14_siglip_gap_896.pali_pt', 'vit_so400m_patch14_siglip_gap_896.pali_refcoco_seg', 'vit_so400m_patch16_siglip_256', 'vit_so400m_patch16_siglip_256.v2_webli', 'vit_so400m_patch16_siglip_256.webli_i18n', 'vit_so400m_patch16_siglip_384', 'vit_so400m_patch16_siglip_384.v2_webli', 'vit_so400m_patch16_siglip_512', 'vit_so400m_patch16_siglip_512.v2_webli', 'vit_so400m_patch16_siglip_gap_256', 'vit_so400m_patch16_siglip_gap_256.v2_webli', 'vit_so400m_patch16_siglip_gap_256.webli_i18n', 'vit_so400m_patch16_siglip_gap_384', 'vit_so400m_patch16_siglip_gap_384.v2_webli', 'vit_so400m_patch16_siglip_gap_512', 'vit_so400m_patch16_siglip_gap_512.v2_webli', 'vit_tiny_patch16_224', 'vit_tiny_patch16_224.augreg_in21k', 'vit_tiny_patch16_224.augreg_in21k_ft_in1k', 'vit_tiny_patch16_384', 'vit_tiny_patch16_384.augreg_in21k_ft_in1k', 'vit_wee_patch16_reg1_gap_256', 'vit_wee_patch16_reg1_gap_256.sbb_in1k', 'vit_xsmall_patch16_clip_224', 'vit_xsmall_patch16_clip_224.tinyclip_yfcc15m']) – Name of the ViT architecture to instantiate (Hugging Face models).

  • pretrained (bool) – If set to True, loads pretrained model weights during initialization, else trains a new model (random initialization).

  • img_size (int) – Input image size.

  • drop_path_rate (float) – Stochastic depth drop rate, acts like dropout but at the block level.

  • drop_rate (float) – Classifier (head) dropout rate.

  • global_pool (Literal['', 'avg', 'avgmax', 'max', 'token', 'map']) – Type of global pooling for final sequence (default: ‘token’).

  • peft_config (PeftConfig | None) – The configuration object (PeftConfig) containing the parameters of the Peft model. (i.e. LoraConfig).

  • num_classes (int) – Number of outputs (classes) for this model.

property num_classes: int

Number of outputs (classes) for this model.

Returns:

The number of outputs supported by this model.

Return type:

int

on_save_checkpoint(checkpoint)[source]

Perform actions during checkpoint saving (called by lightning).

Called by Lightning when saving a checkpoint to give you a chance to store anything else you might want to save. Use on_load_checkpoint() to restore what additional data is saved here.

Parameters:

checkpoint (MutableMapping[str, Any]) – The checkpoint to save.

on_load_checkpoint(checkpoint)[source]

Perform actions during model loading (called by lightning).

If you saved something with on_save_checkpoint() this is your chance to restore this.

Parameters:

checkpoint (MutableMapping[str, Any]) – The loaded checkpoint.

Return type:

None

forward(x)[source]

Same as torch.nn.Module.forward().

Parameters:
  • *args – Whatever you decide to pass into the forward method.

  • **kwargs – Keyword arguments are also possible.

Returns:

Your model’s output