| |
| import functools |
|
|
| import torch |
|
|
|
|
| def assert_tensor_type(func): |
|
|
| @functools.wraps(func) |
| def wrapper(*args, **kwargs): |
| if not isinstance(args[0].data, torch.Tensor): |
| raise AttributeError( |
| f'{args[0].__class__.__name__} has no attribute ' |
| f'{func.__name__} for type {args[0].datatype}') |
| return func(*args, **kwargs) |
|
|
| return wrapper |
|
|
|
|
| class DataContainer: |
| """A container for any type of objects. |
| |
| Typically tensors will be stacked in the collate function and sliced along |
| some dimension in the scatter function. This behavior has some limitations. |
| 1. All tensors have to be the same size. |
| 2. Types are limited (numpy array or Tensor). |
| |
| We design `DataContainer` and `MMDataParallel` to overcome these |
| limitations. The behavior can be either of the following. |
| |
| - copy to GPU, pad all tensors to the same size and stack them |
| - copy to GPU without stacking |
| - leave the objects as is and pass it to the model |
| - pad_dims specifies the number of last few dimensions to do padding |
| """ |
|
|
| def __init__(self, |
| data, |
| stack=False, |
| padding_value=0, |
| cpu_only=False, |
| pad_dims=2): |
| self._data = data |
| self._cpu_only = cpu_only |
| self._stack = stack |
| self._padding_value = padding_value |
| assert pad_dims in [None, 1, 2, 3] |
| self._pad_dims = pad_dims |
|
|
| def __repr__(self): |
| return f'{self.__class__.__name__}({repr(self.data)})' |
|
|
| def __len__(self): |
| return len(self._data) |
|
|
| @property |
| def data(self): |
| return self._data |
|
|
| @property |
| def datatype(self): |
| if isinstance(self.data, torch.Tensor): |
| return self.data.type() |
| else: |
| return type(self.data) |
|
|
| @property |
| def cpu_only(self): |
| return self._cpu_only |
|
|
| @property |
| def stack(self): |
| return self._stack |
|
|
| @property |
| def padding_value(self): |
| return self._padding_value |
|
|
| @property |
| def pad_dims(self): |
| return self._pad_dims |
|
|
| @assert_tensor_type |
| def size(self, *args, **kwargs): |
| return self.data.size(*args, **kwargs) |
|
|
| @assert_tensor_type |
| def dim(self): |
| return self.data.dim() |
|
|
