| import torch |
| from ._utils import _type, _cuda, _range |
| |
| |
| class _StorageBase(object): |
| is_cuda = False |
| |
| def __str__(self): |
| content = ' ' + '\n '.join(str(self[i]) for i in _range(len(self))) |
| return content + '\n[{} of size {}]'.format(torch.typename(self), len(self)) |
| |
| def __repr__(self): |
| return str(self) |
| |
| def __iter__(self): |
| return iter(map(lambda i: self[i], _range(self.size()))) |
| |
| def copy_(self, source, async=False): |
| if async: |
| torch._C._storageCopyAsync(self, source) |
| else: |
| torch._C._storageCopy(self, source) |
| return self |
| |
| def __copy__(self): |
| return self.clone() |
| |
| def __deepcopy__(self, memo): |
| memo = memo.setdefault('torch', {}) |
| if self._cdata in memo: |
| return memo[self._cdata] |
| new_storage = self.clone() |
| memo[self._cdata] = new_storage |
| return new_storage |
| |
| def __reduce__(self): |
| return type(self), (self.tolist(),) |
| |
| def clone(self): |
| return type(self)(self.size()).copy_(self) |
| |
| def tolist(self): |
| return [v for v in self] |
| |
| def cpu(self): |
| return self.type(getattr(torch, self.__class__.__name__)) |
| |
| def double(self, async=False): |
| return self.type(type(self).__module__ + '.DoubleStorage', async) |
| |
| def float(self, async=False): |
| return self.type(type(self).__module__ + '.FloatStorage', async) |
| |
| def long(self, async=False): |
| return self.type(type(self).__module__ + '.LongStorage', async) |
| |
| def int(self, async=False): |
| return self.type(type(self).__module__ + '.IntStorage', async) |
| |
| def short(self, async=False): |
| return self.type(type(self).__module__ + '.ShortStorage', async) |
| |
| def char(self, async=False): |
| return self.type(type(self).__module__ + '.CharStorage', async) |
| |
| def byte(self, async=False): |
| return self.type(type(self).__module__ + '.ByteStorage', async) |
| |
| def pin_memory(self): |
| if self.is_cuda: |
| raise TypeError("cannot pin '{0}' only CPU memory can be pinned" |
| .format(self.type())) |
| import torch.cuda |
| allocator = torch.cuda._host_allocator() |
| return type(self)(self.size(), allocator=allocator).copy_(self) |
| |
| |
| _StorageBase.type = _type |
| _StorageBase.cuda = _cuda |