Source code for block_zoo.Pooling

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT license.

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np

from block_zoo.BaseLayer import BaseLayer, BaseConf
from utils.DocInherit import DocInherit


class PoolingConf(BaseConf):
    """

    Args:
        pool_type (str): 'max' or 'mean', default is 'max'.
        pool_axis (int): which axis to conduct pooling, default is 1.
    """
    def __init__(self, **kwargs):
        super(PoolingConf, self).__init__(**kwargs)

    @DocInherit
    def default(self):
        #self.input_dim = 128
        self.pool_type = 'max'  # Supported: ['max', mean']
        self.pool_axis = 1

    @DocInherit
    def declare(self):
        self.num_of_inputs = 1
        self.input_ranks = [-1]

    @DocInherit
    def inference(self):
        self.output_dim = []
        for idx, dim in enumerate(self.input_dims[0]):
            if idx != self.pool_axis:
                self.output_dim.append(dim)

        # DON'T MODIFY THIS
        self.output_rank = len(self.output_dim)

    @DocInherit
    def verify(self):
        super(PoolingConf, self).verify()

        necessary_attrs_for_user = ['pool_type', 'pool_axis']
        for attr in necessary_attrs_for_user:
            self.add_attr_exist_assertion_for_user(attr)

        self.add_attr_value_assertion('pool_type', ['max', 'mean'])

        assert all([input_rank >= 3 for input_rank in self.input_ranks]), "Cannot apply a pooling layer on a tensor of which the rank is less than 3. Usually, a tensor whose rank is at least 3, e.g. [batch size, sequence length, feature]"

        assert self.output_dim[-1] != -1, "Pooling on the axis %d while the input shape is %s requires that the sequence lengths should be fixed! Please set it on conf/training_params/fixed_lengths" % (self.pool_axis, str(self.input_dims[0]))

class Pooling(BaseLayer):
    """ Pooling layer

    Args:
        layer_conf (PoolingConf): configuration of a layer
    """
    def __init__(self, layer_conf):
        super(Pooling, self).__init__(layer_conf)

    def forward(self, string, string_len=None):
        """ process inputs

        Args:
            string (Tensor): any shape.
            string_len (Tensor): [batch_size], default is None.

        Returns:
            Tensor: Pooling result of string

        """
        if self.layer_conf.pool_type == "mean":
            assert not string_len is None, "Parameter string_len should not be None!"
            string = torch.sum(string, self.layer_conf.pool_axis).squeeze(self.layer_conf.pool_axis)
            if not torch.is_tensor(string_len):
                string_len = torch.FloatTensor(string_len).unsqueeze(1)
            if self.is_cuda():
                device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
                string_len = string_len.to(device)
            string_len = string_len.unsqueeze(1)
            output = string / string_len.expand_as(string).float()
        elif self.layer_conf.pool_type == "max":
            output = torch.max(string, self.layer_conf.pool_axis)[0]

        return output, string_len