grid_compression_head

Bases: head

A head for compressing grid data using geometric compression.

Supports different patch structures (cuboid, cylinder, sphere) and packing strategies for grid data.

Attributes:

Name	Type	Description
h	int	Height of the grid.
w	int	Width of the grid.
channel_num	int	Number of channels in the grid.
d	int, default=1	Depth of the grid.
name	str	Name of the head.
pooling_metric	str, default='batch_max'	Metric used for pooling operations.
patch_shape	str, default='cuboid'	Shape of the patch. Options: 'cuboid', 'cylinder', 'sphere'.
cd_h, cd_w, cd_d	(int, optional)	Parameters for packing and compression.
packing_strategy	str, default='densest_packing'	Strategy for packing grid patches.
with_dropout	bool, default=True	Whether to apply dropout during output processing.
p	float, default=0.5	Dropout probability.
parameters_init_method	str, default='xavier_normal'	Initialization method for parameters.
device	str, default='cpu'	Device to run the computations ('cpu' or 'cuda').

Methods:

Name	Description
get_patch_size	Returns the size of the patch.
get_input_grid_shape	Returns the shape of the input grid.
get_output_grid_shape	Returns the shape of the output grid after packing.

Source code in tinybig/head/grid_based_heads.py

class grid_compression_head(head):
    """
    A head for compressing grid data using geometric compression.

    Supports different patch structures (cuboid, cylinder, sphere) and packing strategies for grid data.

    Attributes
    ----------
    h : int
        Height of the grid.
    w : int
        Width of the grid.
    channel_num : int
        Number of channels in the grid.
    d : int, default=1
        Depth of the grid.
    name : str
        Name of the head.
    pooling_metric : str, default='batch_max'
        Metric used for pooling operations.
    patch_shape : str, default='cuboid'
        Shape of the patch. Options: 'cuboid', 'cylinder', 'sphere'.
    cd_h, cd_w, cd_d : int, optional
        Parameters for packing and compression.
    packing_strategy : str, default='densest_packing'
        Strategy for packing grid patches.
    with_dropout : bool, default=True
        Whether to apply dropout during output processing.
    p : float, default=0.5
        Dropout probability.
    parameters_init_method : str, default='xavier_normal'
        Initialization method for parameters.
    device : str, default='cpu'
        Device to run the computations ('cpu' or 'cuda').

    Methods
    -------
    get_patch_size()
        Returns the size of the patch.
    get_input_grid_shape()
        Returns the shape of the input grid.
    get_output_grid_shape()
        Returns the shape of the output grid after packing.
    """
    def __init__(
        self,
        h: int, w: int, channel_num: int,
        d: int = 1, name: str = 'grid_compression_head',
        pooling_metric: str = 'batch_max',
        patch_shape: str = 'cuboid',
        p_h: int = None, p_h_prime: int = None,
        p_w: int = None, p_w_prime: int = None,
        p_d: int = 0, p_d_prime: int = None,
        p_r: int = None,
        cd_h: int = None, cd_w: int = None, cd_d: int = 1,
        packing_strategy: str = 'densest_packing',
        with_dropout: bool = True, p: float = 0.5,
        # other parameters
        parameters_init_method: str = 'xavier_normal',
        device: str = 'cpu', *args, **kwargs
    ):
        """
        Initializes the `grid_compression_head` class.

        Parameters
        ----------
        h : int
            Height of the grid.
        w : int
            Width of the grid.
        channel_num : int
            Number of channels in the grid.
        d : int, default=1
            Depth of the grid.
        name : str, default='grid_compression_head'
            Name of the head.
        pooling_metric : str, default='batch_max'
            Metric to use for pooling operations.
        patch_shape : str, default='cuboid'
            Shape of the patch. Options: 'cuboid', 'cylinder', 'sphere'.
        p_h : int, optional
            Patch height.
        p_h_prime : int, optional
            Adjusted patch height for the cuboid.
        p_w : int, optional
            Patch width. Defaults to `p_h` if not provided.
        p_w_prime : int, optional
            Adjusted patch width for the cuboid.
        p_d : int, default=0
            Patch depth.
        p_d_prime : int, optional
            Adjusted patch depth for the cuboid.
        p_r : int, optional
            Patch radius (for spherical or cylindrical patches).
        cd_h : int, optional
            Compression depth in the height dimension.
        cd_w : int, optional
            Compression depth in the width dimension.
        cd_d : int, default=1
            Compression depth in the depth dimension.
        packing_strategy : str, default='densest_packing'
            Strategy for packing patches into the grid.
        with_dropout : bool, default=True
            Whether to apply dropout during output processing.
        p : float, default=0.5
            Dropout probability.
        parameters_init_method : str, default='xavier_normal'
            Initialization method for model parameters.
        device : str, default='cpu'
            Device to run the computations ('cpu' or 'cuda').
        """
        if channel_num is None or channel_num <=0:
            raise ValueError(f'positive channel number={channel_num} must be specified...')
        self.channel_num = channel_num
        if h is None or w is None or d is None:
            raise ValueError(f'h={h} and w={w} and d={d} must be specified...')
        grid_structure = grid(
            h=h, w=w, d=d, universe_num=channel_num
        )

        if patch_shape == 'cuboid':
            assert p_h is not None
            p_w = p_w if p_w is not None else p_h
            patch_structure = cuboid(p_h=p_h, p_w=p_w, p_d=p_d, p_h_prime=p_h_prime, p_w_prime=p_w_prime, p_d_prime=p_d_prime)
        elif patch_shape == 'cylinder':
            assert p_r is not None
            patch_structure = cylinder(p_r=p_r, p_d=p_d, p_d_prime=p_d_prime)
        elif patch_shape == 'sphere':
            assert p_r is not None
            patch_structure = sphere(p_r=p_r)
        else:
            raise ValueError(f'patch_shape={patch_shape} must be either cuboid, cylinder or sphere...')

        data_transformation = geometric_compression(
            grid=grid_structure,
            patch=patch_structure,
            packing_strategy=packing_strategy,
            cd_h=cd_h, cd_w=cd_w, cd_d=cd_d,
            metric=partial(metric, metric_name=pooling_metric),
            device=device,
        )

        remainder = zero_remainder(
            device=device,
        )

        output_process_functions = []
        if with_dropout:
            output_process_functions.append(torch.nn.Dropout(p=p))
        print('pooling layer', output_process_functions)

        m = data_transformation.get_grid_size(across_universe=True)
        n = data_transformation.get_patch_num(across_universe=True)

        super().__init__(
            m=m, n=n,
            name=name,
            data_transformation=data_transformation,
            remainder=remainder,
            output_process_functions=output_process_functions,
            parameters_init_method=parameters_init_method,
            device=device, *args, **kwargs
        )

    def get_patch_size(self):
        """
        Returns the size of the patch in the grid.

        Returns
        -------
        int
            Patch size.
        """
        return self.data_transformation.get_patch_size()

    def get_input_grid_shape(self):
        """
        Returns the shape of the input grid.

        Returns
        -------
        tuple
            A tuple representing (height, width, depth) of the input grid.
        """
        return self.data_transformation.get_grid_shape()

    def get_output_grid_shape(self):
        """
        Returns the shape of the output grid after packing.

        Returns
        -------
        tuple
            A tuple representing (height, width, depth) of the packed grid.
        """
        output_h, output_w, output_d = self.data_transformation.get_grid_shape_after_packing()
        return output_h, output_w, output_d

__init__(h, w, channel_num, d=1, name='grid_compression_head', pooling_metric='batch_max', patch_shape='cuboid', p_h=None, p_h_prime=None, p_w=None, p_w_prime=None, p_d=0, p_d_prime=None, p_r=None, cd_h=None, cd_w=None, cd_d=1, packing_strategy='densest_packing', with_dropout=True, p=0.5, parameters_init_method='xavier_normal', device='cpu', *args, **kwargs)

Initializes the grid_compression_head class.

Parameters:

Name	Type	Description	Default
h	int	Height of the grid.	required
w	int	Width of the grid.	required
channel_num	int	Number of channels in the grid.	required
d	int	Depth of the grid.	1
name	str	Name of the head.	'grid_compression_head'
pooling_metric	str	Metric to use for pooling operations.	'batch_max'
patch_shape	str	Shape of the patch. Options: 'cuboid', 'cylinder', 'sphere'.	'cuboid'
p_h	int	Patch height.	None
p_h_prime	int	Adjusted patch height for the cuboid.	None
p_w	int	Patch width. Defaults to p_h if not provided.	None
p_w_prime	int	Adjusted patch width for the cuboid.	None
p_d	int	Patch depth.	0
p_d_prime	int	Adjusted patch depth for the cuboid.	None
p_r	int	Patch radius (for spherical or cylindrical patches).	None
cd_h	int	Compression depth in the height dimension.	None
cd_w	int	Compression depth in the width dimension.	None
cd_d	int	Compression depth in the depth dimension.	1
packing_strategy	str	Strategy for packing patches into the grid.	'densest_packing'
with_dropout	bool	Whether to apply dropout during output processing.	True
p	float	Dropout probability.	0.5
parameters_init_method	str	Initialization method for model parameters.	'xavier_normal'
device	str	Device to run the computations ('cpu' or 'cuda').	'cpu'

Source code in tinybig/head/grid_based_heads.py

def __init__(
    self,
    h: int, w: int, channel_num: int,
    d: int = 1, name: str = 'grid_compression_head',
    pooling_metric: str = 'batch_max',
    patch_shape: str = 'cuboid',
    p_h: int = None, p_h_prime: int = None,
    p_w: int = None, p_w_prime: int = None,
    p_d: int = 0, p_d_prime: int = None,
    p_r: int = None,
    cd_h: int = None, cd_w: int = None, cd_d: int = 1,
    packing_strategy: str = 'densest_packing',
    with_dropout: bool = True, p: float = 0.5,
    # other parameters
    parameters_init_method: str = 'xavier_normal',
    device: str = 'cpu', *args, **kwargs
):
    """
    Initializes the `grid_compression_head` class.

    Parameters
    ----------
    h : int
        Height of the grid.
    w : int
        Width of the grid.
    channel_num : int
        Number of channels in the grid.
    d : int, default=1
        Depth of the grid.
    name : str, default='grid_compression_head'
        Name of the head.
    pooling_metric : str, default='batch_max'
        Metric to use for pooling operations.
    patch_shape : str, default='cuboid'
        Shape of the patch. Options: 'cuboid', 'cylinder', 'sphere'.
    p_h : int, optional
        Patch height.
    p_h_prime : int, optional
        Adjusted patch height for the cuboid.
    p_w : int, optional
        Patch width. Defaults to `p_h` if not provided.
    p_w_prime : int, optional
        Adjusted patch width for the cuboid.
    p_d : int, default=0
        Patch depth.
    p_d_prime : int, optional
        Adjusted patch depth for the cuboid.
    p_r : int, optional
        Patch radius (for spherical or cylindrical patches).
    cd_h : int, optional
        Compression depth in the height dimension.
    cd_w : int, optional
        Compression depth in the width dimension.
    cd_d : int, default=1
        Compression depth in the depth dimension.
    packing_strategy : str, default='densest_packing'
        Strategy for packing patches into the grid.
    with_dropout : bool, default=True
        Whether to apply dropout during output processing.
    p : float, default=0.5
        Dropout probability.
    parameters_init_method : str, default='xavier_normal'
        Initialization method for model parameters.
    device : str, default='cpu'
        Device to run the computations ('cpu' or 'cuda').
    """
    if channel_num is None or channel_num <=0:
        raise ValueError(f'positive channel number={channel_num} must be specified...')
    self.channel_num = channel_num
    if h is None or w is None or d is None:
        raise ValueError(f'h={h} and w={w} and d={d} must be specified...')
    grid_structure = grid(
        h=h, w=w, d=d, universe_num=channel_num
    )

    if patch_shape == 'cuboid':
        assert p_h is not None
        p_w = p_w if p_w is not None else p_h
        patch_structure = cuboid(p_h=p_h, p_w=p_w, p_d=p_d, p_h_prime=p_h_prime, p_w_prime=p_w_prime, p_d_prime=p_d_prime)
    elif patch_shape == 'cylinder':
        assert p_r is not None
        patch_structure = cylinder(p_r=p_r, p_d=p_d, p_d_prime=p_d_prime)
    elif patch_shape == 'sphere':
        assert p_r is not None
        patch_structure = sphere(p_r=p_r)
    else:
        raise ValueError(f'patch_shape={patch_shape} must be either cuboid, cylinder or sphere...')

    data_transformation = geometric_compression(
        grid=grid_structure,
        patch=patch_structure,
        packing_strategy=packing_strategy,
        cd_h=cd_h, cd_w=cd_w, cd_d=cd_d,
        metric=partial(metric, metric_name=pooling_metric),
        device=device,
    )

    remainder = zero_remainder(
        device=device,
    )

    output_process_functions = []
    if with_dropout:
        output_process_functions.append(torch.nn.Dropout(p=p))
    print('pooling layer', output_process_functions)

    m = data_transformation.get_grid_size(across_universe=True)
    n = data_transformation.get_patch_num(across_universe=True)

    super().__init__(
        m=m, n=n,
        name=name,
        data_transformation=data_transformation,
        remainder=remainder,
        output_process_functions=output_process_functions,
        parameters_init_method=parameters_init_method,
        device=device, *args, **kwargs
    )

get_input_grid_shape()

Returns the shape of the input grid.

Returns:

Type	Description
tuple	A tuple representing (height, width, depth) of the input grid.

Source code in tinybig/head/grid_based_heads.py

def get_input_grid_shape(self):
    """
    Returns the shape of the input grid.

    Returns
    -------
    tuple
        A tuple representing (height, width, depth) of the input grid.
    """
    return self.data_transformation.get_grid_shape()

get_output_grid_shape()

Returns the shape of the output grid after packing.

Returns:

Type	Description
tuple	A tuple representing (height, width, depth) of the packed grid.

Source code in tinybig/head/grid_based_heads.py

def get_output_grid_shape(self):
    """
    Returns the shape of the output grid after packing.

    Returns
    -------
    tuple
        A tuple representing (height, width, depth) of the packed grid.
    """
    output_h, output_w, output_d = self.data_transformation.get_grid_shape_after_packing()
    return output_h, output_w, output_d

get_patch_size()

Returns the size of the patch in the grid.

Returns:

Type	Description
int	Patch size.

Source code in tinybig/head/grid_based_heads.py

def get_patch_size(self):
    """
    Returns the size of the patch in the grid.

    Returns
    -------
    int
        Patch size.
    """
    return self.data_transformation.get_patch_size()