Skip to content

special_function_process

Special function processing method.

It handles some special functions to accommodate their requirements, like the batchnorms.

Parameters:

Name Type Description Default
func_class

The function class information.

 required
func_parameters

The dictionary of function parameters.

 required
device

The device for hosting and processing these special functions.

 'cpu'

Returns:

Type Description
tuple | list

The tuple of processed function class, and function parameters.
Source code in tinybig/util/util.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def special_function_process(func_class, func_parameters, device='cpu', *args, **kwargs):
    """
    Special function processing method.

    It handles some special functions to accommodate their requirements, like the batchnorms.

    Parameters
    ----------
    func_class: str
        The function class information.
    func_parameters: dict
        The dictionary of function parameters.
    device: str, default = 'cpu'
        The device for hosting and processing these special functions.

    Returns
    -------
    tuple | list
        The tuple of processed function class, and function parameters.
    """
    if func_class in ['torch.nn.BatchNorm1d', 'torch.nn.BatchNorm2d', 'torch.nn.BatchNorm3d']:
        if 'device' not in func_parameters:
            func_parameters['device'] = device
    return func_class, func_parameters