comfy.py

"""
    "PyTorch must serve man, not rule over him" (c) DAO

A module for simple conversion between numpy and torch types.
Created out of frustration from lines like:
 - x = Variable(torch.FloatTensor(x)).cuda()                 # now var(x, 'float32')
 - (model(x).max(1)[1].data.cpu().numpy() == y).mean()       # now numpy(x)
"""

import torch
from torch.autograd import Variable
import numpy as np
import scipy.sparse
from warnings import warn


def var(x, dtype=None, device=None, sparse=None,
        requires_grad=None, volatile=None, detach=False,
        **kwargs):
    """
    Does everything in its power to cast x to a variable.
    :param x: something to be converted.
        Supports lists, tuples, numpy arrays, torch tensors, variables
    :param dtype: numpy dtype; if specified, casts x to this dtype. Preserves x type by default.
    :param device: can be "cpu", "gpu0", "gpu1", ..., "gpuN", (in gpu memory)
        0,1, ..., N (alias for "gpu0", ... )
        default is torch default device
    :param sparse: by default respects input sparsity
        (i.e. scipy sparse matrix will become sparse variable, dense np.array will become dense variable)
        if True, returns a variable with torch sparse tensor even if x is "dense"
        if False, returns a variable with "dense" (normal) tensor even if x is a sparse matrix
    :param requires_grad: boolean indicating whether the Variable has been
        created by a subgraph containing any Variable, that requires it.
        See :ref:`excluding-subgraphs` for more details.
        Can be changed only on leaf Variables.
    :param volatile: boolean indicating that the Variable should be used in
        inference mode, i.e. don't save the history. See
        :ref:`excluding-subgraphs` for more details.
        Can be changed only on leaf Variables.
    :param detach: if True, creates a new leaf variable even if x is already Variable.
        Useful to avoid gradient propagation through x.
    :param kwargs: any additional params used when casting x to a numpy array (intermediate step)
        Examples: copy=False, order='K', ndmin=2, ...
        Used only if x is neither numpy array nor torch object.
    """
    if isinstance(x, Variable):
        if detach:
            x = x.detach()
        x.requires_grad = requires_grad or x.requires_grad
        x.volatile = volatile or x.volatile

        if dtype is None:
            return x
        else:
            if is_sparse(x) and not sparse:
                x = x.clone()  # create new child variable (non-inplace)
                x.data = x.data.to_dense()
                if x.grad is not None:
                    x.grad = x.grad.to_dense()
            elif not is_sparse(x) and sparse:
                if not detach:
                    warn("var: conversion of variable \n%s\n from dense to sparse "
                         "will not propagate gradients!" % x)
                x = Variable(to_sparse(x.data),
                             requires_grad=requires_grad or False,
                             volatile=volatile or False)
            return x.type(torch_type(dtype, device, sparse))

    else:  # type(x) in torch tensor, np array, list, scipy sparse matrix, any custom iterable
        return Variable(tensor(x, dtype, device, sparse, **kwargs),
                        requires_grad=requires_grad or False,
                        volatile=volatile or False)


def tensor(x, dtype=None, device=None, sparse=None, **kwargs):
    """
    Does everything in its power to convert x to some torch tensor.
    Works like np.array but for torchies.
    :param x: something to be converted.
        Supports lists, tuples, numpy arrays, torch tensors, variables, etc.
    :param dtype: numpy dtype; if specified, casts x to this dtype. Preserves x type by default.
    :param sparse: by default, converts scipy.sparse to sparse tensors, others to "dense"
        if True, returns torch sparse tensor for any input.
        if False, returns "dense" (normal) tensors even if x is a sparse matrix
    :param device: can be "cpu", "gpu0", "gpu1", ..., "gpuN", (in gpu memory)
        0,1, ..., N (alias for "gpu0", ... )
        default is torch default device
    :param kwargs: any additional params used when casting x to a numpy array (intermediate step)
        Examples: copy=False, order='K', ndmin=2, ...
        Used only if x is neither numpy array nor torch object.

    """
    def check_type(x):
        assert torch.is_tensor(x)
        if is_sparse(x) and not sparse:
            x = x.to_dense()
        elif not is_sparse(x) and sparse:
            x = to_sparse(x)
        return x.type(torch_type(dtype or numpy_dtype(type(x)), device, sparse))

    if isinstance(x, Variable):
        return check_type(x.data)
    elif torch.is_tensor(x):
        return check_type(x)
    elif scipy.sparse.issparse(x):
        return check_type(scipy_sparse_to_torch(x, device=device))
    elif is_numpy_object(x):
        return check_type(torch.from_numpy(x))
    else:  # type(x) in list, tuple, nested tuple, custom iterable, pandas series, etc.
        x = np.array(x, dtype=dtype, **kwargs)
        return check_type(torch.from_numpy(x))


def torch_type(dtype, device=None, sparse=False):
    """
    Returns a torch class that corresponds to given numpy dtype.
    :param dtype: string or numpy dtype.
    :param sparse: boolean, True means torch.sparse, False is otherwise ("dense")
    :param device: can be "cpu", "gpu" (default gpu),
        "gpu0", "gpu1", ..., "gpuN", (in gpu memory)
        0,1, ..., N (alias for "gpu0", ... )
        default is torch default device
    """
    # find CPU type
    base_type = type(torch.from_numpy(np.array([0], np.dtype(dtype))))
    if sparse:
        type_fullname = torch.typename(base_type)
        assert type_fullname.startswith("torch.") and type_fullname.endswith("Tensor") \
               and type_fullname.count('.') >= 1, "Bad type name: %s" % type_fullname

        typename = type_fullname.split('.')[-1]  # "*Tensor"
        assert hasattr(torch.sparse, typename), "There is no torch.sparse type for %s" % type_fullname
        base_type = getattr(torch.sparse, typename)

    # handle device
    assert isinstance(device, int) or device in (None, 'cpu', b'cpu') or device.startswith('gpu'), \
        "Bad device name : %s" % device
    if device is None:
        device = get_default_device()

    if device == 'cpu':
        return base_type

    else:  # device = gpu# or just integer

        # infer device index from string
        if isinstance(device, str) or isinstance(device, bytes):
            device = str(device)
            if device == 'gpu':
                device = None  # default cuda
            else:
                assert device.startswith('gpu')
                try:
                    device = int(device[3:])
                except:
                    raise IndexError("Cannot infer device index from device: %s " % device)

        dummy_obj = base_type([])
        return type(dummy_obj.cuda(device))


def numpy_dtype(torch_type):
    """
    return numpy dtype corresponding to torch type
    """
    type_fullname = torch.typename(torch_type)
    assert type_fullname.startswith('torch') and type_fullname.endswith('Tensor')
    type_code = type_fullname.split('.')[-1][:-len('Tensor')]
    type_mapping = {
        'Byte': 'uint8',
        'Char': 'uint8',
        'Double': 'float64',
        'Float': 'float32',
        'Half': 'float16',
        'Int': 'int32',
        'Long': 'int64',
        'Short': 'int16'
    }
    return np.dtype(type_mapping.get(type_code))


def is_sparse(x):
    """ returns True if x is torch sparse tensor, False if "dense" tensor """
    typename = torch.typename(type(x))
    return 'sparse' in typename


def is_on_gpu(x):
    typename = torch.typename(type(x))
    return 'cuda' in typename


def to_sparse(x):
    """ converts dense tensor x to sparse format """
    assert not isinstance(x, Variable), "to_sparse does not support Variable at this point" \
                                        "because of issues with gradiet propagation (know how2fix? contribute!)"
    sparse_tensortype = torch_type(numpy_dtype(type(x)), sparse=True)
    indices = torch.nonzero(x)
    if len(indices.shape) == 0:  # if all elements are zeros
        return sparse_tensortype(*x.shape)
    indices = indices.t()
    values = x[tuple(indices[i] for i in range(indices.shape[0]))]
    x_sparse = sparse_tensortype(indices.cpu(), values.cpu(), x.size())
    return x_sparse


def scipy_sparse_to_torch(x, device=None):
    """ converts scipy sparse matrix to torch sparse tensor """
    assert scipy.sparse.issparse(x), "this function only supports scipy sparse matrices"
    sparse_type = torch_type(x.dtype, device=device, sparse=True)
    indices_type = torch_type('int64', device=device, sparse=False)
    values_type = torch_type(x.dtype, device=device, sparse=False)
    i, j, values = scipy.sparse.find(x)
    indices = indices_type(np.stack([i, j]))
    values = values_type(values)
    return sparse_type(indices, values, torch.Size(x.shape))


def get_default_device():
    if torch.cuda.is_available():
        if 'cuda' in torch.typename(torch.zeros(0)):
            return torch.cuda.current_device()
    return 'cpu'


def is_numpy_object(x):
    """checks if var is a numpy array"""
    return type(x).__module__.startswith("numpy")


def numpy(x, dtype=None, order=None):
    """ Does all in it's power to convert x from torch types to a numpy array """
    if isinstance(x, Variable):
        x = x.data
    if is_sparse(x):
        x = x.to_dense()
    if torch.is_tensor(x):
        x = x.cpu().numpy()
    return np.asarray(x, dtype=dtype, order=order)

torch_comfy.ipynb