File size: 1,307 Bytes
ddadf19 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
# coding: utf-8
__author__ = 'cleardusk'
import os
import numpy as np
import torch
import pickle
def mkdir(d):
os.makedirs(d, exist_ok=True)
def _get_suffix(filename):
"""a.jpg -> jpg"""
pos = filename.rfind('.')
if pos == -1:
return ''
return filename[pos + 1:]
def _load(fp):
suffix = _get_suffix(fp)
if suffix == 'npy':
return np.load(fp)
elif suffix == 'pkl':
return pickle.load(open(fp, 'rb'))
def _dump(wfp, obj):
suffix = _get_suffix(wfp)
if suffix == 'npy':
np.save(wfp, obj)
elif suffix == 'pkl':
pickle.dump(obj, open(wfp, 'wb'))
else:
raise Exception('Unknown Type: {}'.format(suffix))
def _load_tensor(fp, mode='cpu'):
if mode.lower() == 'cpu':
return torch.from_numpy(_load(fp))
elif mode.lower() == 'gpu':
return torch.from_numpy(_load(fp)).cuda()
def _tensor_to_cuda(x):
if x.is_cuda:
return x
else:
return x.cuda()
def _load_gpu(fp):
return torch.from_numpy(_load(fp)).cuda()
_load_cpu = _load
_numpy_to_tensor = lambda x: torch.from_numpy(x)
_tensor_to_numpy = lambda x: x.numpy()
_numpy_to_cuda = lambda x: _tensor_to_cuda(torch.from_numpy(x))
_cuda_to_tensor = lambda x: x.cpu()
_cuda_to_numpy = lambda x: x.cpu().numpy()
|