def _set_loa(self, loa):
self._loa = loa
uvg = UniqueValueGetter(
'LOARandomSampleDataFlow length consistency check failed')
for i in self._loa:
uvg.set(len(i))
self._length = uvg.get()
def __init__(self, loa):
self._loa = loa
uvg = UniqueValueGetter(
'ListOfArrayDataFlow length consistency check failed')
for i in self._loa:
uvg.set(len(i))
self._length = uvg.get()
def _stack(self, values, process):
uvg = UniqueValueGetter('Tensor sizes should match except the first dim.')
for v in values:
uvg.set(v.size()[1:])
uvg.get()
out = None
if torchdl._use_shared_memory:
# If we're in a background process, concatenate directly into a
# shared memory tensor to avoid an extra copy
numel = 0
if process:
if self._mode is VarLengthCollateMode.PAD:
numel = max([x.numel() for x in values]) * len(values)
elif self._mode is VarLengthCollateMode.CONCAT:
numel = sum([x.numel() for x in values])
elif self._mode is VarLengthCollateMode.PAD2D:
max_h = max([x.size(0) for x in values])
max_w = max([x.size(1) for x in values])
hw = max_h * max_w
numel = sum([x.numel() // x.size(0) // x.size(1) * hw for x in values])
else:
numel = sum([x.numel() for x in values])
if numel > 0:
storage = values[0].storage()._new_shared(numel)
out = values[0].new(storage)
if not process:
return torch.stack(values, 0, out=out)
lengths = [v.size(0) for v in values]
if self._mode is VarLengthCollateMode.CONCAT:
return torch.cat(values, 0, out=out), torch.LongTensor(lengths)
elif self._mode is VarLengthCollateMode.PAD:
max_length = max(lengths)
result = []
for v in values:
if v.size(0) < max_length:
v = torch.cat([v, v.new(*((max_length - v.size(0), ) + v.size()[1:])).zero_()], dim=0)
result.append(v)
return torch.stack(result, 0, out=out), torch.LongTensor(lengths)
else:
raise ValueError('Unknown collation mode: {}.'.format(self._mode))
def _stack(self, values, key=None):
mode, parameters = None, None
if key is not None:
mode_spec = self._fields[key]
if isinstance(mode_spec, tuple):
mode = VarLengthCollateMode.from_string(mode_spec[0])
parameters = mode_spec[1:]
else:
mode = VarLengthCollateMode.from_string(mode_spec)
parameters = tuple()
out = None
if use_shared_memory():
# If we're in a background process, concatenate directly into a
# shared memory tensor to avoid an extra copy
numel = 0
if key is not None:
if mode is VarLengthCollateMode.PAD:
numel = max([x.numel() for x in values]) * len(values)
elif mode is VarLengthCollateMode.CONCAT:
numel = sum([x.numel() for x in values])
elif mode is VarLengthCollateMode.PAD2D:
max_h = max([x.size(0) for x in values])
max_w = max([x.size(1) for x in values])
hw = max_h * max_w
numel = sum([
x.numel() // x.size(0) // x.size(1) * hw
for x in values
])
else:
numel = sum([x.numel() for x in values])
if numel > 0:
storage = values[0].storage()._new_shared(numel)
out = values[0].new(storage)
if key is None:
return torch.stack(values, 0, out=out)
if mode is VarLengthCollateMode.CONCAT:
uvg = UniqueValueGetter(
'Tensor sizes should match except the first dim.')
for v in values:
uvg.set(v.size()[1:])
lengths = [v.size(0) for v in values]
return torch.cat(values, 0, out=out), torch.LongTensor(lengths)
elif mode is VarLengthCollateMode.PAD:
uvg = UniqueValueGetter(
'Tensor sizes should match except the first dim.')
for v in values:
uvg.set(v.size()[1:])
pad_value = parameters[0] if len(parameters) > 0 else 0
lengths = [v.size(0) for v in values]
max_length = max(lengths)
result = []
for v in values:
if v.size(0) < max_length:
v = torch.cat([
v,
v.new(*((max_length - v.size(0), ) +
v.size()[1:])).fill_(pad_value)
],
dim=0)
result.append(v)
return torch.stack(result, 0, out=out), torch.LongTensor(lengths)
elif mode is VarLengthCollateMode.PAD2D:
uvg = UniqueValueGetter(
'Tensor sizes should match except the first 2 dims.')
for v in values:
uvg.set(v.size()[2:])
rest_size = uvg.get() or []
pad_value = parameters[0] if len(parameters) > 0 else 0
lengths = [v.size()[:2] for v in values]
max_h, max_w = max([x[0] for x in lengths
]), max([x[1] for x in lengths])
result = []
for v in values:
u = v.new(*(max_h, max_w, *rest_size)).fill_(pad_value)
u[:v.size(0), :v.size(1)] = v
result.append(u)
return torch.stack(result, 0, out=out), torch.LongTensor(lengths)
elif mode is VarLengthCollateMode.PADIMAGE:
uvg = UniqueValueGetter(
'Tensor sizes should match except the last 2 dims.')
for v in values:
assert v.dim() == 3, 'Support only 3-dimensional input.'
uvg.set(v.size(0))
pad_value = parameters[0] if len(parameters) > 0 else 0
lengths = [v.size()[-2:] for v in values]
max_h, max_w = max([x[0] for x in lengths
]), max([x[1] for x in lengths])
result = []
for v in values:
u = v.new(*(uvg.get(), max_h, max_w)).fill_(pad_value)
# TODO(Jiayuan Mao @ 07/19): support input with dim > 3.
u[:, :v.size(1), :v.size(2)] = v
result.append(u)
return torch.stack(result, 0, out=out), torch.LongTensor(lengths)
else:
raise ValueError('Unknown collation mode: {}.'.format(mode))
def _stack(self, values, spec=None, maybe_cuda=True):
mode = spec.type if spec is not None else None
out = None
if self.mode is VarLengthCollateV3Mode.COLLATE and use_shared_memory():
# If we're in a background process, concatenate directly into a
# shared memory tensor to avoid an extra copy
numel = 0
if mode is not None:
if mode is DataLayoutType.CONCAT:
numel = sum([x.numel() for x in values])
elif mode is DataLayoutType.PAD:
numel = max([x.numel() for x in values]) * len(values)
elif mode is DataLayoutType.PAD2D:
max_h = max([x.size(0) for x in values])
max_w = max([x.size(1) for x in values])
hw = max_h * max_w
numel = sum([
x.numel() // x.size(0) // x.size(1) * hw
for x in values
])
elif mode is DataLayoutType.PADIMAGE:
max_h = max([x.size(1) for x in values])
max_w = max([x.size(2) for x in values])
hw = max_h * max_w
numel = sum([
x.numel() // x.size(1) // x.size(2) * hw
for x in values
])
else:
numel = sum([x.numel() for x in values])
if numel > 0:
storage = values[0].storage()._new_shared(numel)
out = values[0].new(storage)
if mode is None:
return self._stack_raw(values, out=out, maybe_cuda=maybe_cuda)
if mode is DataLayoutType.CONCAT:
uvg = UniqueValueGetter(
'Tensor sizes should match except the first dim.')
for v in values:
uvg.set(v.size()[1:])
lengths = [v.size(0) for v in values]
return _VarLengthCollateV3ArrayStack(
self._stack_raw(result, out=out, maybe_cuda=maybe_cuda),
torch.LongTensor(lengths))
elif mode is DataLayoutType.PAD:
uvg = UniqueValueGetter(
'Tensor sizes should match except the first dim.')
for v in values:
uvg.set(v.size()[1:])
pad_value = spec.fill
lengths = [v.size(0) for v in values]
max_length = max(lengths)
result = []
for v in values:
if v.size(0) < max_length:
v = torch.cat([
v,
v.new(*((max_length - v.size(0), ) +
v.size()[1:])).fill_(pad_value)
],
dim=0)
result.append(v)
return _VarLengthCollateV3ArrayStack(
self._stack_raw(result, out=out, maybe_cuda=maybe_cuda),
torch.LongTensor(lengths))
elif mode is DataLayoutType.PAD2D:
uvg = UniqueValueGetter(
'Tensor sizes should match except the first 2 dims.')
for v in values:
uvg.set(v.size()[2:])
rest_size = uvg.get() or []
pad_value = spec.fill
lengths = [v.size()[:2] for v in values]
max_h, max_w = max([x[0] for x in lengths
]), max([x[1] for x in lengths])
result = []
for v in values:
u = v.new(*(max_h, max_w, *rest_size)).fill_(pad_value)
u[:v.size(0), :v.size(1)] = v
result.append(u)
return _VarLengthCollateV3ArrayStack(
self._stack_raw(result, out=out, maybe_cuda=maybe_cuda),
torch.LongTensor(lengths))
elif mode is DataLayoutType.PADIMAGE:
uvg = UniqueValueGetter(
'Tensor sizes should match except the last 2 dims.')
for v in values:
assert v.dim() == 3, 'Support only 3-dimensional input.'
uvg.set(v.size(0))
pad_value = spec.fill
lengths = [v.size()[-2:] for v in values]
max_h, max_w = max([x[0] for x in lengths
]), max([x[1] for x in lengths])
result = []
for v in values:
u = v.new(*(uvg.get(), max_h, max_w)).fill_(pad_value)
# TODO(Jiayuan Mao @ 07/19): support input with dim > 3.
u[:, :v.size(1), :v.size(2)] = v
result.append(u)
return _VarLengthCollateV3ArrayStack(
self._stack_raw(result, out=out, maybe_cuda=maybe_cuda),
torch.LongTensor(lengths))
else:
raise ValueError('Unknown collation mode: {}.'.format(mode))