Python API#

Provides the easy-to-use fine-grain algorithmic control over an existing DNN

The framework currently features two methods for algorithmic swapping. swap_backend() which swaps every module type of a DNN returning an object completely managed by ai3 and swap_conv2d() which swaps convolution operations out of the existing DNN.

Functions Providing Algorithmic Selection#

ai3.swap_conv2d(module, algos: str | Sequence[str] | Callable | None = None, sample_input_shape: Sequence[int] | None = None)#

Swaps, in-place, conv2d operations out of the existing DNN for an implementation of the user specified algorithm. After swapping, the same DNN can still be trained and compiled. If no AlgorithmicSelector is given then the default algorithm decided by the framework are used.

Parameters:

  • module – the module containing the conv2d operations to be swapped out in-place

  • algos (Optional[AlgorithmicSelector]) – algorithmic selector for the conv2d operations

  • sample_input_shape – the input shape to the DNN which is passed to the function algorithmic selector if present

Example

Swaps the first conv2d operation for an implementation of direct convolution and the second conv2d operation for an implementation of SMM convolution

>>> input_data = torch.randn(10, 3, 224, 224)
>>> orig = ConvNet()
>>> orig_out = orig(input_data)
>>> ai3.swap_conv2d(orig, ['direct', 'smm'])
>>> sc_out = orig(input_data)
>>> torch.allclose(orig_out, sc_out, atol=1e-6)
True
ai3.swap_backend(module, algos: Mapping[str, str | Sequence[str] | Callable] | None = None, sample_input_shape: Sequence[int] | None = None, *, dtype=None) → Model#

Swaps every module in an exsiting DNN for an implementation of the user specified algorithm returning a Model completly managed by the framework.

Algorithmic selection is performed by passing a mapping from strings containing names of the operations to swap to a AlgorithmicSelector. If no AlgorithmicSelector is passed for a given operation then the default algorithm decided by the framework are used.

Parameters:

  • module – the module to have its operations replaced

  • algos (Optional[Mapping[str, AlgorithmicSelector]]) – mapping from operation type to a AlgorithmicSelector

  • sample_input_shape – the input shape to the DNN which is passed to the function algorithmic selector if present

Returns:

A Model which uses the algorithms specified and yields equivalent outputs as the original DNN

Example: Swaps the first conv2d operation for an implementation of direct convolution and the second conv2d operation for an implementation of SMM convolution

>>> def auto_selector(orig: torch.nn.Conv2d, input_shape) -> str:
...     out_channels = orig.weight.shape[0]
...     if (out_channels < 50 and
...         input_shape[1] < 50 and
...         input_shape[2] > 150 and
...         input_shape[3] > 150):
...         return 'direct'
...     return 'smm'
...
>>> input_data = torch.randn(1, 3, 224, 224)
>>> vgg16 = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights.DEFAULT)
>>> vgg16 = vgg16.eval()
>>> with torch.inference_mode():
...     torch_out = vgg16(input_data)
...     model: ai3.Model = ai3.swap_backend(vgg16, {"conv2d": auto_selector,
...                                                 "maxpool2d": "default"},
...                                         sample_input_shape=(1, 3, 224, 224))
...     sb_out = model(input_data)
...     torch.allclose(torch_out, sb_out, atol=1e-4)
True

Types#

ai3.AlgorithmicSelector#

The object that performs the algorithmic selection for an associated operation

There are three different types of algorithmic selectors.

  • str: contains the name of the algorithm that will be used for all modules

    of the associated type

  • Sequence[str]: list of algorithm names, as modules are encountered in a forward pass,

    they are replaced with an implementation of the algorithm in the list with the same index as that module has relative to the other modules of the associated type

  • Callable: function which is given the module being swapped and returns the name of the

    algorithm to use, the function can optionally take the input size for this module provided that a sample input shape was passed to the swapping function.

Example

Function to perform algorithmic selection.

>>> def conv2d_selector(orig: torch.nn.Conv2d, input_shape) -> str:
...     out_channels = orig.weight.shape[0]
...     if (out_channels < 50 and
...         input_shape[0] < 50 and
...         input_shape[1] > 150 and
...         input_shape[2] > 150):
...         return 'direct'
...     return 'smm'
...
>>> input_data = torch.randn(10, 3, 224, 224)
>>> orig = ConvNet()
>>> torch_out = orig(input_data)
>>> ai3.swap_conv2d(orig, conv2d_selector, (3, 224, 224))
>>> sc_out = orig(input_data)
>>> torch.allclose(torch_out, sc_out, atol=1e-6)
True

alias of str | Sequence[str] | Callable

Classes#

class ai3.Model(layers: Sequence[Layer])#

The model which performs the operations using the user specified algorithms.

predict(input, out_type=None)#

Returns the output after passing input through the layers.

Parameters:

  • input – input tensor to perform operations on

  • out_type – the desired type of the output tensor

Returns:

output after performing operations

class ai3.Tensor(tens: Tensor)#

Simple type which implements the Python Buffer Protocol

numpy()#

Transform this Tensor to a numpy.ndarray using the buffer protocol.

Returns:

numpy.ndarray with the shape and contents of the original

to(out_type)#

Transform this Tensor to another type using the buffer protocol.

Parameters:

out_type – type to transform this object to

Returns:

An object of out_type with the shape and contents of the original

torch()#

Transform this Tensor to a torch.Tensor using the buffer protocol.

Returns:

torch.Tensor with the shape and contents of the original

Global Variables#

ai3.SUPPORTED_ALGORITHMS#

The supported operations and their supported algorithms.

See Supported Operations, their Algorithms, and Acceleration Platform Compatibility for supported acceleration platform by algorithm.

{ 'conv2d': [ 'direct',
              'smm',
              'winograd',
              'gemm',
              'implicit gemm',
              'implicit precomp gemm',
              'guess',
              'mps',
              'metal',
              'default'],
  'linear': ['direct', 'default'],
  'maxpool2d': ['direct', 'default'],
  'avgpool2d': ['direct', 'default'],
  'adaptiveavgpool2d': ['direct', 'default'],
  'relu': ['direct', 'default'],
  'flatten': ['direct', 'default']}
ai3.FROM_BACKEND: str | None = None#

The backend of the existing DNN

If None, the backend will be inferred from the type of the input to the swapping functions

ai3.SUPPORTED_FROM_BACKENDS: Sequence[str] = ['torch']#

Algorithmic selection over a DNN is supported for these backends

Other Functions#

ai3.using_sycl() → bool#

Whether the implementations can use SYCL

ai3.using_cudnn() → bool#

Whether the implementations can use cuDNN

ai3.using_cublas() → bool#

Whether the implementations can use cuBLAS

ai3.using_mps_and_metal() → bool#

Whether the implementations can use MPS and Metal