| | |
| | 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() |
| |
|
