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

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.datasets import load_digits\n",
    "import torch, torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "X,y = load_digits(return_X_y=True)\n",
    "\n",
    "model = nn.Sequential()\n",
    "model.add_module('dense1', nn.Linear(64, 50))\n",
    "model.add_module('relu1', nn.ReLU())\n",
    "model.add_module('dense2', nn.Linear(50, 10))\n",
    "model.cuda(0)\n",
    "\n",
    "opt = torch.optim.RMSprop(model.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Pure pytorch__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss =  3.1297069\n",
      "acc =  0.0859375\n",
      "loss =  32.540497\n",
      "acc =  0.1171875\n",
      "loss =  27.85704\n",
      "acc =  0.328125\n",
      "loss =  15.00075\n",
      "acc =  0.1015625\n",
      "loss =  14.075615\n",
      "acc =  0.2578125\n",
      "loss =  5.1595225\n",
      "acc =  0.2421875\n",
      "loss =  3.074504\n",
      "acc =  0.3671875\n",
      "loss =  1.3797708\n",
      "acc =  0.5234375\n",
      "loss =  1.2865769\n",
      "acc =  0.546875\n",
      "loss =  1.0730145\n",
      "acc =  0.5703125\n"
     ]
    }
   ],
   "source": [
    "from torch.autograd import Variable\n",
    "for i in range(10):\n",
    "    batch = np.random.randint(0, len(X), 128)\n",
    "    X_batch = Variable(torch.FloatTensor(X[batch])).cuda(0)\n",
    "    y_batch = Variable(torch.LongTensor(y[batch])).cuda(0)\n",
    "    y_pred = model(X_batch)\n",
    "    loss = F.cross_entropy(y_pred, y_batch)\n",
    "    \n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    opt.zero_grad()\n",
    "    \n",
    "    print('loss = ', loss.data.cpu().numpy()[0])\n",
    "    print('acc = ', torch.eq(y_pred.max(1)[1], y_batch).data.cpu().numpy().mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__pytorch + comfy__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss = 0.85640407\n",
      "acc = 0.0859375\n",
      "loss = 0.93089414\n",
      "acc = 0.0859375\n",
      "loss = 0.9754122\n",
      "acc = 0.109375\n",
      "loss = 0.72724664\n",
      "acc = 0.140625\n",
      "loss = 0.5673185\n",
      "acc = 0.0546875\n",
      "loss = 0.63406694\n",
      "acc = 0.171875\n",
      "loss = 0.48260045\n",
      "acc = 0.0859375\n",
      "loss = 0.5338222\n",
      "acc = 0.1875\n",
      "loss = 0.40716457\n",
      "acc = 0.171875\n",
      "loss = 0.42137316\n",
      "acc = 0.0625\n"
     ]
    }
   ],
   "source": [
    "from comfy import var, tensor, numpy\n",
    "for i in range(10):\n",
    "    batch = np.random.randint(0, len(X), 128)\n",
    "    y_pred = model(var(X[batch], 'float32', device=0))\n",
    "    loss = F.cross_entropy(y_pred, var(y[batch], device=0))\n",
    "    \n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    opt.zero_grad()\n",
    "    \n",
    "    print('loss =', numpy(loss)[0])\n",
    "    print('acc =', numpy(torch.eq(y_pred.max(1)[1], y_batch)).mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__more stuff__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from comfy import to_sparse, scipy_sparse_to_torch\n",
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "\n",
    "data = [\n",
    "    'there is no emotion , there is peace',\n",
    "    'there is no ignorence , there is knowledge',\n",
    "    'there is no passion , there is serenity',\n",
    "    'there is no chaos , there is harmony',\n",
    "    'there is no death , there is the force'\n",
    "]\n",
    "\n",
    "data = CountVectorizer(dtype='float32').fit_transform(data)\n",
    "\n",
    "data = scipy_sparse_to_torch(data)\n",
    "weights = torch.randn(data.shape[1], 16)\n",
    "\n",
    "emb = data.matmul(weights)  # sparse dot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}