Torchvision models

Model providing helper funtions to create TorchVision Deep Learning models.

TorchVision pages:

Import

from ikomia.dnn.torch import models

Functions

faster_rcnn([train_mode, use_pretrained, ...])

Create Torchvision Faster RCNN model for training or inference.

mask_rcnn([train_mode, use_pretrained, ...])

Create Torchvision Mask RCNN model for training or inference.

mnasnet([train_mode, use_pretrained, ...])

Create Torchvision MNasNet model for training or inference.

resnet([model_name, train_mode, ...])

Create Torchvision ResNet model for training or inference.

resnext([model_name, train_mode, ...])

Create Torchvision ResNeXt model for training or inference.

Details

ikomia.dnn.torch.models.faster_rcnn(train_mode: bool = False, use_pretrained: bool = True, input_size: int = 800, classes: int = 2)

Create Torchvision Faster RCNN model for training or inference.

Parameters:

  • train_mode (bool) – True or False

  • use_pretrained (bool) – True to do transfer learning from pre-trained model, False to train from scratch

  • input_size (int) – input image size

  • classes (int) – number of classes in the dataset

Returns:

model object

ikomia.dnn.torch.models.mask_rcnn(train_mode: bool = False, use_pretrained: bool = True, input_size: int = 800, classes: int = 2)

Create Torchvision Mask RCNN model for training or inference.

Parameters:

  • train_mode (boolean) – True or False

  • use_pretrained (boolean) – True to do transfer learning from pre-trained model, False to train from scratch

  • input_size (int) – input image size

  • classes (int) – number of classes in the dataset

Returns:

model object

ikomia.dnn.torch.models.mnasnet(train_mode: bool = False, use_pretrained: bool = False, feature_extract: bool = False, classes: int = 2)

Create Torchvision MNasNet model for training or inference.

Parameters:

  • train_mode (bool) – True or False

  • use_pretrained (bool) – True to do transfer learning from pre-trained model, False to train from scratch

  • feature_extract (bool) – transfer learning only, True to keep pre-trained features (train last layers only), False to train all layers

  • classes (int) – number of classes in the dataset

Returns:

model object

ikomia.dnn.torch.models.resnet(model_name: str = 'resnet50', train_mode: bool = False, use_pretrained: bool = False, feature_extract: bool = False, classes: int = 2)

Create Torchvision ResNet model for training or inference.

Parameters:

  • model_name (str) – model name

  • train_mode (bool) – True or False

  • use_pretrained (bool) – True to do transfer learning from pre-trained model, False to train from scratch

  • feature_extract (bool) – transfer learning only, True to keep pre-trained features (train last layers only), False to train all layers

  • classes (int) – number of classes in the dataset

Returns:

model object

ikomia.dnn.torch.models.resnext(model_name: str = 'resnext50', train_mode: bool = False, use_pretrained: bool = False, feature_extract: bool = False, classes: int = 2)

Create Torchvision ResNeXt model for training or inference.

Parameters:

  • model_name (str) – model name

  • train_mode (bool) – True or False

  • use_pretrained (bool) – True to do transfer learning from pre-trained model, False to train from scratch

  • feature_extract (bool) – transfer learning only, True to keep pre-trained features (train last layers only), False to train all layers

  • classes (int) – number of classes in the dataset

Returns:

model object