ai4med.components.models package
NVIDIA Docs Hub  NVIDIA Clara  NVIDIA Clara Train 3.1  ai4med.components.models package

Submodules

ai4med.components.models.alexnet module

class Alexnet(num_classes, num_feat='96,256,512,1024,1024,3072,4096', mil=False, use_batch_norm=False, use_group_norm=True, use_group_normG=8, reg_weight=0.0, dropout_prob=0.5, final_activation='softmax', dtype=tf.float32, data_format='channels_first')

Bases: ai4med.components.models.model.Model

WZ. 9/10/2019 Build AlexNet for 3D images or 2D images on image classification with / without (fixed input size) multiple instance learning (MIL)

Parameters

  • inputs (tensor) – Tensor of input image (e.g NxCxHxWxD)

  • num_classes (int) – number of classes

  • training (bool) – bool variable indicating training or testing

  • num_feat (int) – number of filters or hidden nodes in each layer

  • mil (bool) – use multi-instance learning or not

  • use_batch_norm (bool) – bool variable indicating batch norm, default True

  • use_group_norm (bool) – bool variable indicating group norm, default False

  • use_group_normG (int) – number of groups

  • data_format (str) – channels_first (default) or channels_last

Returns

The output of the AlexNet (MIL)

get_loss()

Get the additional loss function in AHNet model.

Args:

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the AHNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

ai4med.components.models.densenet121 module

class DenseNet121(weight_decay, pretrain_weight_name=None, pretrain_weight_url=None, data_format='channels_last', n_spatial_dim=2)

Bases: ai4med.components.models.model.Model

Classification model based on DenseNet-121.

Keras densenet121 model with intial weights taken from: https://github.com/flyyufelix/DenseNet-Keras

URL to the h5 file: https://drive.google.com/a/nvidia.com/uc?authuser=0&id=0Byy2AcGyEVxfSTA4SHJVOHNuTXc&export=download

Dense Convolutional Network (DenseNet) connects each layer to every other layer in a feed-forward fashion. Whereas traditional convolutional networks with L layers have L connections - one between each layer and its subsequent layer - DenseNet has L(L+1)/2 direct connections. For each layer, the feature maps of all preceding layers are used as inputs, and its own feature maps are used as inputs into all subsequent layers. It can support both 2D and 3D classification tasks.

For more implementation details of DenseNet121, please check the libs.

Parameters

  • weight_decay (float) – L2 regularization factor

  • pretrain_weight_name (str) – local path of pretrain weight file

  • pretrain_weight_url (str) – link to download pretrain weight file

  • n_spatial_dim (int) – spatial dimension count of data, support 2D and 3D data

Returns

Prediction results

get_loss()

Get the additional loss function in DenseNet121 model.

Args:

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of DenseNet121 model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the AHNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

get_update_ops()

Get the update_ops for Batch Normalization.

The method “tf.control_dependencies” allow the operations used as inputs of the context manager are run before the operations defined inside the context manager. So we use “update_ops” to implement Batch Normalization.

Args:

Returns

Update operations

ai4med.components.models.model module

class Model

Bases: ai4med.common.graph_component.GraphComponent

Base class of Models

Args:

Returns

Prediction results

build(build_ctx: ai4med.common.build_ctx.BuildContext)

Connect model with graph.

Parameters

build_ctx – specified graph context

Returns

Prediction results

get_loss()

Get the additional loss function in AHNet model.

Args:

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the AHNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

get_update_ops()

Get the update_ops for Batch Normalization.

The method “tf.control_dependencies” allow the operations used as inputs of the context manager are run before the operations defined inside the context manager. So we use “update_ops” to implement Batch Normalization.

Args:

Returns

Update operations

ai4med.components.models.seg_ahnet module

class SegAhnet(num_classes, if_use_psp, plane, pretrain_weight_name, pretrain_weight_url=None, final_activation='softmax', dtype=tf.float32, data_format='channels_first', n_spatial_dim=2)

Bases: ai4med.components.models.model.Model

Segmentation model based on AHNet.

3D Anisotropic Hybrid Network (AH-Net) transfers convolutional features learned from 2D images to 3D anisotropic volumes. Such a transfer inherits the desired strong generalization capability for within-slice information while naturally exploiting between-slice information for more effective modelling. It can support both 2D and 3D segmentation tasks.

This implementation accepts an hdf5 file as pretrain weights through pretrain_weight_url, and if provided will initialize the network with those weights. If nothing is provided, the default initial weights will be from: https://github.com/fchollet/deep-learning-models/releases/download/v0.2/resnet50_weights_tf_dim_ordering_tf_kernels.h5

For more implementation details of AHNet, please check the libs.

Note: The input dimensions must be divisible 32 due to the structure of layers in the network.

Parameters

  • num_classes (int) – expected class number of prediction results

  • if_use_psp (bool) – use “pyramid volumetric pooling” module for decoder or not

  • plane (str) – permute the weight tensors of first convolution layer to change channel dimension to plane dimension

  • pretrain_weight_name (str) – local path of pretrain weight file

  • pretrain_weight_url (str) – link to download pretrain weight file

  • final_activation (str) – final activation layer, supports: softmax, sigmoid, linear

  • dtype (tf.dtype) – reserved argument for future features

  • data_format (str) – input and output data dimensions supports CHANNEL_FIRST and CHANNEL_LAST

  • n_spatial_dim (int) – spatial dimension count of data, support 2D and 3D data

Returns

Prediction results

get_loss()

Get the additional loss function in AHNet model.

Args:

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of AHNet model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the AHNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

ai4med.components.models.seg_resnet module

class SegResnet(num_classes, blocks_down='1,2,2,4', blocks_up='1,1,1', init_filters=8, use_batch_norm=False, use_group_norm=True, use_group_normG=8, reg_weight=0.0, dropout_prob=0.0, final_activation='softmax', use_vae=False, dtype=tf.float32, data_format='channels_first')

Bases: ai4med.components.models.model.Model

3D Segmentation model based on ResNet encoder and decoder.

Construct a encoder - decoder structure network based on ResNet.

Parameters

  • num_classes (int) – expected class number of prediction results

  • blocks_down (str) – encoder layer number and ResNet block number of every layer

  • blocks_up (str) – decoder layer number and ResNet block number of every layer

  • init_filters (int) – convolution kernel number of first layer

  • use_batch_norm (bool) – use batch normalization, but only if use_group_norm is explicitly set to false.

  • use_group_norm (bool) – use group normalization or not. By default this is true.

  • use_group_normG (int) – group numbter of group normalization

  • reg_weight (float) – weight regularization factor

  • dropout_prob (float) – dropout rate for several layers in training

  • final_activation (str) – final activation layer, supports: softmax, sigmoid, linear

  • use_vae (bool) – use “Variational Auto Encoder” or not

  • dtype (tf.dtype) – data type of the network

  • data_format (str) – input and output data dimensions supports CHANNEL_FIRST and CHANNEL_LAST

Returns

Prediction results

get_loss()

Get the additional loss function in AHNet model.

Parameters

None –

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of ResNet model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the ResNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

ai4med.components.models.unet module

class Unet(num_classes, nf_enc='32,64,64,64', nf_dec='64,64,64,64,64,32,32', use_batch_norm=False, use_group_norm=True, use_group_normG=8, reg_weight=0.0, dropout_prob=0.0, final_activation='softmax', dtype=tf.float32, data_format='channels_first')

Bases: ai4med.components.models.model.Model

3D Segmentation model based on UNet.

For more details, check our implemented UNet in the paper. Wentao Zhu et al., NeurReg: Neural Registration and Its Application to Image Segmentation, WACV 2020. Wentao Zhu et al., Neural Multi-Scale Self-Supervised Registration for Echocardiogram Dense Tracking, arXiv 2019.

Parameters

  • num_classes (int) – expected class number of prediction results

  • nf_enc (str) – number of filters in each encoder layer

  • nf_dec (str) – number of filters in each decoder layer

  • use_batch_norm (bool) – augmented batch normalization

  • use_group_norm (bool) – augmented group normalization

  • use_group_normG (int) – group numbter of group normalization

  • reg_weight (float) – weight regularization factor

  • dropout_prob (float) – dropout rate for several layers in training

  • final_activation (str) – final activation layer, supports: softmax, sigmoid, linear

  • dtype (tf.dtype) – data type of the network

  • data_format (str) – input and output data dimensions supports CHANNEL_FIRST and CHANNEL_LAST

Returns

Prediction results

get_loss()

Get the additional loss function in AHNet model.

Args:

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the AHNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

ai4med.components.models.unet_parallel module

class UnetParallel(num_classes, nf_enc='32,64,64,64', nf_dec='64,64,64,64,64,32,32', gpu_enc='0,0,0,0', gpu_dec='0,0,0,0,0,0,0', use_batch_norm=False, use_group_norm=True, use_group_normG=8, reg_weight=0.0, dropout_prob=0.0, final_activation='softmax', dtype=tf.float32, data_format='channels_first')

Bases: ai4med.components.models.model.Model

3D Segmentation model based on UNet.

For more details, check our implemented UNet in the paper. Wentao Zhu et al., NeurReg: Neural Registration and Its Application to Image Segmentation, WACV 2020. Wentao Zhu et al., Neural Multi-Scale Self-Supervised Registration for Echocardiogram Dense Tracking, arXiv 2019.

Parameters

  • num_classes (int) – expected class number of prediction results

  • nf_enc (str) – number of filters in each encoder layer

  • nf_dec (str) – number of filters in each decoder layer

  • gpu_enc (str) – the gpu id for each encoder layer

  • gpu_dec (str) – the gpu id for each decoder layer

  • use_batch_norm (bool) – augmented batch normalization

  • use_group_norm (bool) – augmented group normalization

  • use_group_normG (int) – group numbter of group normalization

  • reg_weight (float) – weight regularization factor

  • dropout_prob (float) – dropout rate for several layers in training

  • final_activation (str) – final activation layer, supports: softmax, sigmoid, linear

  • dtype (tf.dtype) – data type of the network

  • data_format (str) – input and output data dimensions supports CHANNEL_FIRST and CHANNEL_LAST

Returns

Prediction results

get_loss()

Get the additional loss function in AHNet model.

Args:

Returns

Loss function

get_predictions(inputs, is_training, build_ctx: ai4med.common.build_ctx.BuildContext)

Forward computation process of model for both training and inference.

Parameters

  • inputs (tf.Tensor) – input data for the AHNet model

  • is_training (bool) – in training process or not

  • build_ctx (BuildContext) – reserved argument for future features

Returns

Prediction results

Module contents

© Copyright 2020, NVIDIA. Last updated on Feb 2, 2023.