torch_utils.py

import numpy as np
import torch, os
import torch.nn.utils.rnn as rnn_utils
from typing import Tuple
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import torchvision
from torchvision import transforms


def flatten(x):
    '''
    flatten high dimensional tensor x into an array
    :param x: shape (B, D1, D2, ...)
    :return: 1 dimensional tensor
    '''
    dims = x.size()[1:] #remove the first dimension as it is batch dimension
    num_features = 1
    for s in dims: num_features *= s
    return x.contiguous().view(-1, num_features)


def gpu(tensor, gpu=False):

    if gpu: return tensor.cuda()
    else: return tensor


def cpu(tensor):
    if tensor.is_cuda: return tensor.cpu()
    else: return tensor


def minibatch(*tensors, **kwargs):

    batch_size = kwargs['batch_size']

    if len(tensors) == 1:
        tensor = tensors[0]
        for i in range(0, len(tensor), batch_size):
            yield tensor[i:i + batch_size]
    else:
        for i in range(0, len(tensors[0]), batch_size):
            yield tuple(x[i:i + batch_size] for x in tensors)


def shuffle(*arrays, **kwargs):
    """This is not an inplace operation. Therefore, you can shuffle without worrying changing data."""
    if len(set(len(x) for x in arrays)) != 1:
        raise ValueError('All inputs to shuffle must have '
                         'the same length.')

    shuffle_indices = np.arange(len(arrays[0]))
    np.random.shuffle(shuffle_indices) # fix this for reproducible

    if len(arrays) == 1:
        return arrays[0][shuffle_indices]
    else:
        return tuple(x[shuffle_indices] for x in arrays)


def assert_no_grad(variable):

    if variable.requires_grad:
        raise ValueError(
            "nn criterions don't compute the gradient w.r.t. targets - please "
            "mark these variables as volatile or not requiring gradients"
        )


def numpy2tensor(x, dtype):
    # torch.tensor(torch.from_numpy(var), dtype = torch.int, torch.long)
    return torch.tensor(torch.from_numpy(x), dtype = dtype)


def tensor2numpy(x):
    # return x.numpy()
    return cpu(x).numpy()


def set_seed(seed, cuda=False):

    torch.manual_seed(seed)
    if cuda: torch.cuda.manual_seed(seed)


def create_mask_tensor(query: torch.Tensor, doc: torch.Tensor, threshold: int = 0):
    """
    Creating masking of two tensor. These two tensors are integer tensor
    Parameters

    ----------
    query: (B, L)
    doc: (B, R)
    threshold: when it is 0, means we ignore padding tokens. when it is 1, it means we ignore <unk> or oov words
    Returns
    -------

    """
    assert query.size(0) == doc.size(0)
    assert len(query.size()) == 2 and len(doc.size()) == 2
    query_mask = query > threshold
    doc_mask = doc > threshold
    query_mask = query_mask.unsqueeze(2)  # (B, L, 1)
    doc_mask = doc_mask.unsqueeze(2)  # (B, R, 1)
    doc_mask = doc_mask.permute(0, 2, 1)  # (B, 1, R)

    mask_tensor = torch.bmm(query_mask.float(), doc_mask.float())  # (B, L, R)
    return mask_tensor  # , torch.sum(query_mask, dim = 1).squeeze(), torch.sum(doc_mask, dim = 1).squeeze()


def create_mask_tensor_image(left_indices: torch.Tensor, right_indices: torch.Tensor, threshold: int = 0):
    """
    Creating masking of two tensor. These two tensors are integer tensor
    Parameters

    ----------
    left_indices: (B1, n1, M1)
    right_indices: (B, n, M2)
    threshold: when it is 0, means we ignore padding tokens. when it is 1, it means we ignore <unk> or oov words
    Returns
    -------

    """
    B1, n1, M1 = left_indices.size()
    B, n, M2 = right_indices.size()
    assert n1 == 1
    left_mask = left_indices > 0
    right_mask = right_indices > 0
    left_mask = left_mask.view(B1, M1, 1)
    if B1 == 1: left_mask = left_mask.expand(B, M1, 1)  # during testing
    right_mask = right_mask.view(B, n * M2, 1)
    ans = torch.bmm(left_mask.float(), right_mask.permute(0, 2, 1).float())
    ans = ans.view(B, M1, n, M2).permute(0, 2, 1, 3)  # (B, n, M1, M2)
    return ans


def count_parameters(model: nn.Module):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def get_sorted_index_and_reverse_index(base_array: np.ndarray):
    """
    We use sorted_index = np.argsort(-base_array) to find the indices to short the array decreasingly.
    We also need to find the indices to restoring the original order of elements of base_array
    after apply sorted_index.
    This method is important because we need to input the tensor to GRU/LSTM with packed sequence.
    Parameters
    ----------
    base_array: (B, )

    Returns
    -------

    """
    assert type(base_array) == np.ndarray
    batch_size = base_array.shape[0]
    assert base_array.shape == (batch_size,)
    new_indices = np.argsort(-base_array)
    old_indices = np.arange(batch_size)
    r = np.stack([new_indices, old_indices], axis = 1)
    r = r[np.argsort(r[:, 0])]
    restoring_indices = r[:, 1]  # the retoring indices. This method is tested very carefully.
    return new_indices, restoring_indices


def packing_sequence(seq: torch.Tensor, seq_lens: np.ndarray, new_index) -> torch.Tensor:
    """
    Prepare a packed sequence to input to an RNN. It is required that the length of sequences in `seq` must be sorted.
    After

    Parameters
    ----------
    seq: (B, L, D) where L is length of sequence
    seq_lens: (B, )
    new_index: (B, ) this index is used to make sequence lengths sorted
    old_index: (B, ) this index is used to restore the sequence lengths
    Returns
    -------

    """
    return rnn_utils.pack_padded_sequence(seq[new_index], seq_lens[new_index], batch_first = True)


def torch_repeat_dim0(A: torch.tensor, n: int):
    """
    Repeat tensor across a dimension
    Parameters
    ----------
    A
    axis

    Returns
    -------

    """
    assert len(A.size()) == 3
    d1, d2, d3 = A.size()
    A = A.unsqueeze(0).transpose(0, 1).repeat(1, n, 1, 1).view(-1, d2, d3)
    assert A.size() == (n * d1, d2, d3)
    return A


def boolean_mask(target: torch.Tensor, mask: torch.Tensor):
    """
    Mimick tf.boolean_mask
    Copied from https://discuss.pytorch.org/t/slicing-tensor-using-boolean-list/7354/3
    Parameters
    ----------
    target
    mask

    Returns
    -------

    """
    x = mask == True
    # y=torch.arange(0,3)
    # x=torch.Tensor([True,False,True])==True
    # print(y[x])
    return target[x]

def torch_argsort(input, dim=None, descending=False):
    """Returns the indices that sort a tensor along a given dimension in ascending
        order by value.
         This is the second value returned by :meth:`torch.sort`.  See its documentation
        for the exact semantics of this method.
         Args:
            input (Tensor): the input tensor
            dim (int, optional): the dimension to sort along
            descending (bool, optional): controls the sorting order (ascending or descending)
         Example::
             >>> a = torch.randn(4, 4)
            >>> a
            tensor([[ 0.0785,  1.5267, -0.8521,  0.4065],
                    [ 0.1598,  0.0788, -0.0745, -1.2700],
                    [ 1.2208,  1.0722, -0.7064,  1.2564],
                    [ 0.0669, -0.2318, -0.8229, -0.9280]])
             >>> torch.argsort(a, dim=1)
            tensor([[2, 0, 3, 1],
                    [3, 2, 1, 0],
                    [2, 1, 0, 3],
                    [3, 2, 1, 0]])
    """
    # copy from https://github.com/pytorch/pytorch/pull/9600/files
    if dim is None:
        return torch.sort(input, -1, descending)[1]
    return torch.sort(input, dim, descending)[1]


def _predict_process_ids(user_ids, item_ids, num_items, use_cuda):
    """

    Parameters
    ----------
    user_ids
    item_ids
    num_items
    use_cuda

    Returns
    -------

    """
    if item_ids is None:
        item_ids = np.arange(num_items, dtype=np.int64)

    if np.isscalar(user_ids):
        user_ids = np.array(user_ids, dtype=np.int64)

    user_ids = torch.from_numpy(user_ids.reshape(-1, 1).astype(np.int64))
    item_ids = torch.from_numpy(item_ids.reshape(-1, 1).astype(np.int64))
    if item_ids.size()[0] != user_ids.size(0):
        user_ids = user_ids.expand(item_ids.size())

    user_var = gpu(user_ids, use_cuda)
    item_var = gpu(item_ids, use_cuda)

    return user_var.squeeze(), item_var.squeeze()


def idf(total_docs: int, term_freq: int) -> float:
    """compute inverse doc frequency. If a term appears at all docs, then, its value is low for discrimination.
    If a term does not show in any doc, then, we simply use set denominator to 1 => largest idf value """
    assert term_freq <= total_docs, "The number of documents that contain a term must be smaller than total_docs"
    return np.log((1.0 + total_docs) / float(term_freq + 1.0)) + 1.0


def moving_average(input_tensor: torch.Tensor, window_size: int, dimension: int):
    """

    Parameters
    ----------
    input_tensor: torch.Tensor  of shape (B, L, D)
    window_size: sliding windows size
    dimension: dimension we want to apply sliding window

    Returns
    -------

    """
    ret = torch.cumsum(input_tensor, dim = dimension)
    # print("Here:", ret, ret.shape)
    ret[:, window_size:] = ret[:, window_size:] - ret[:, :-window_size]
    return ret[:, window_size - 1:] / window_size


def cosine_distance(a: torch.Tensor, b: torch.Tensor):
    """
    Compute the cosine distance between two tensors. This implementation saves a lot of memory since
    memory complexity is O(B x L x R)
    Parameters
    ----------
    a: `torch.Tensor` shape (B, L, D)
    b: `torch.Tensor` shape (B, R, D)

    Returns
    -------

    """
    assert len(a.size()) == len(b.size()) == 3
    A_square = (a * a).sum(dim = - 1)  # B, L
    B_square = (b * b).sum(dim = -1)  # B, R
    dot = torch.bmm(a, b.permute(0, 2, 1))  # B, L, R
    # added abs in case of negative, added 1e-10 to avoid nan gradient of sqrt
    return torch.sqrt(torch.abs(A_square.unsqueeze(-1) - 2 * dot + B_square.unsqueeze(1)) + 1e-10)


def l1_distance(a: torch.Tensor, b: torch.Tensor):
    """
    Compute the l1 distance between two tensors. This implementation consumes a lot of memory since
    mem complexity is O(B x L x R x D) due to x - y. I tried many ways but this is the best thing I can do
    Parameters
    ----------
    a: `torch.Tensor` shape (B, L, D)
    b: `torch.Tensor` shape (B, R, D)

    Returns
    -------

    """
    assert len(a.size()) == len(b.size()) == 3
    x = a.unsqueeze(2)  # (B, L, 1, D)
    y = b.unsqueeze(1)  # (B, 1, R, D)
    return torch.norm(x - y, p = 1, dim = -1)


def _get_doc_context_copacrr(doc: torch.Tensor, doc_mask: torch.Tensor, context_window_size: int) -> torch.Tensor:
    """

    Parameters
    ----------
    doc: with shape (B, R, D)
    doc_mask: binary tensor that differentiate real tokens from padding tokens (B, R)

    Returns
    -------
    a tensor of shape (B, R, D) which indicates the context representation of each token in doc.
    We also reset padding tokens to zero since they have no context
    """

    def moving_average(a: torch.Tensor, window_size: int, dimension: int):
        ret = torch.cumsum(a, dim = dimension)
        # print("Here:", ret, ret.shape)
        ret[:, window_size:] = ret[:, window_size:] - ret[:, :-window_size]
        return ret[:, window_size - 1:] / window_size

    left = context_window_size // 2
    right = context_window_size - left - 1  # in case context windows is an even number then left=x//2, right=x-x//2
    y = F.pad(doc, (0, 0, left, right))  # (B, c/2 + R + c/2, D)
    document_context = moving_average(y, window_size = context_window_size, dimension = 1)
    document_context = document_context * doc_mask.unsqueeze(-1).float()
    return document_context


def init_weights(m):
    """
    Copied from https://discuss.pytorch.org/t/how-are-layer-weights-and-biases-initialized-by-default/13073/3
    Examples:
        >>> w = nn.Linear(3, 4)
        >>> w.apply(init_weights)
    """
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
        if hasattr(m.bias, "data"): m.bias.data.fill_(0)
    if isinstance(m, nn.Conv2d):
        torch.nn.init.xavier_uniform_(m.weight)
        if m.bias:
            torch.nn.init.xavier_uniform_(m.bias)


def auto_rnn(rnn_cell: nn.RNN, input_feats: torch.Tensor,
             lens: torch.Tensor, new_indices: torch.Tensor, restoring_indices: torch.Tensor, max_len: int):
    """

    Parameters
    ----------
    rnn_cell : a rnn cell
    input_feats: `torch.Tensor` (B, L, D)
    lens: `torch.Tensor` (B, )
    new_indices: `torch.Tensor` (B, )
    restoring_indices: `torch.Tensor` (B, )
    max_len: int
    Returns
    -------

    """
    return rnn_cell((input_feats, lens, new_indices, restoring_indices), max_len=max_len, return_h=False)[0]


def rnn_forward(rnn_cell, x, return_h=True, max_len=None):
    x, x_len, d_new_indices, d_restoring_indices = x
    # x = self.dropout(x)
    # x_idx = d_new_indices
    x_len_sorted = x_len[d_new_indices]
    # x_len_sorted, x_idx = torch.sort(x_len, descending=True)
    x_sorted = x[d_new_indices]  # x.index_select(dim=0, index=x_idx)
    x_ori_idx = d_restoring_indices
    # _, x_ori_idx = torch.sort(x_idx)

    x_packed = nn.utils.rnn.pack_padded_sequence(x_sorted, x_len_sorted, batch_first=True)
    # x_packed, (h, c) = self.rnn(x_packed)
    x_packed, h = rnn_cell(x_packed)  # this is for GRU not LSTM

    x = nn.utils.rnn.pad_packed_sequence(x_packed, batch_first=True, total_length=max_len)[0]
    # x = x.index_select(dim=0, index=x_ori_idx)
    x = x[x_ori_idx]
    if return_h:
        h = h.permute(1, 0, 2).contiguous().view(-1, h.size(0) * h.size(2))  # .squeeze()
        # h = h.index_select(dim=0, index=x_ori_idx)
        h = h[x_ori_idx]
    return x, h


def rnn_last_h(rnn_cell: nn.RNN, input_feats: torch.Tensor,
               lens: torch.Tensor, new_indices: torch.Tensor, restoring_indices: torch.Tensor, max_len: int):
    """
    return the last hidden vectors of an RNN
    Parameters
    ----------
    rnn_cell : a rnn cell
    input_feats: `torch.Tensor` (B, L, D)
    lens: `torch.Tensor` (B, )
    new_indices: `torch.Tensor` (B, )
    restoring_indices: `torch.Tensor` (B, )
    max_len: int
    Returns
    -------

    """
    return rnn_cell((input_feats, lens, new_indices, restoring_indices), max_len=max_len, return_h=True)[1]


def retrieve_elements_from_indices(tensor: torch.Tensor, indices: torch.Tensor):
    """
    Copied from https://discuss.pytorch.org/t/pooling-using-idices-from-another-max-pooling/37209/4
    How does this work? (Checked
    Parameters
    ----------
    tensor: torch.Tensor shape B, C, L, R
    indices: torch.Tensor shape (B, C, L, R) the values are indices where the last two dimensions are flattened

    Returns
    -------

    """
    flattened_tensor = tensor.flatten(start_dim=2)
    output = flattened_tensor.gather(dim=2, index=indices.flatten(start_dim=2)).view_as(indices)
    return output


data_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


def load_images(infile):
    im = Image.open(infile).convert('RGB')
    return data_transforms(im)