Created
January 22, 2018 06:20
-
-
Save justheuristic/9a3e772f08417a5f339222c412c936e0 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
"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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment