Utils

Deep conversions

argus.utils.deep_to(input: Any, *args, **kwarg) Any[source]

Recursively performs dtype and/or device conversion for tensors and nn modules.

Parameters
  • input – Any input with tensors, tuples, lists, dicts, and other objects.

  • *args – args arguments to torch.Tensor.to().

  • **kwargs – kwargs arguments to torch.Tensor.to().

Returns

Output with converted tensors.

Return type

Any

Example

>>> x = [torch.ones(4, 2, device='cuda:1'),
...      {'target': torch.zeros(4, dtype=torch.uint8)}]
>>> x
[tensor([[1., 1.],
         [1., 1.],
         [1., 1.],
         [1., 1.]], device='cuda:1'),
 {'target': tensor([0, 0, 0, 0], dtype=torch.uint8)}]
>>> deep_to(x, 'cuda:0', dtype=torch.float16)
[tensor([[1., 1.],
         [1., 1.],
         [1., 1.],
         [1., 1.]], device='cuda:0', dtype=torch.float16),
 {'target': tensor([0., 0., 0., 0.], device='cuda:0', dtype=torch.float16)}]
argus.utils.deep_detach(input: Any) Any[source]

Returns new tensors, detached from the current graph without gradient requirement. Recursively performs torch.Tensor.detach().

Parameters

input – Any input with tensors, tuples, lists, dicts, and other objects.

Returns

Output with detached tensors.

Return type

Any

Example

>>> x = [torch.ones(4, 2),
...      {'target': torch.zeros(4, requires_grad=True)}]
>>> x
[tensor([[1., 1.],
         [1., 1.],
         [1., 1.],
         [1., 1.]]),
 {'target': tensor([0., 0., 0., 0.], requires_grad=True)}]
>>> deep_detach(x)
[tensor([[1., 1.],
         [1., 1.],
         [1., 1.],
         [1., 1.]]),
 {'target': tensor([0., 0., 0., 0.])}]
argus.utils.deep_chunk(input: Any, chunks: int, dim: int = 0) List[Any][source]

Slice tensors into approximately equal chunks. Duplicates references to objects that are not tensors. Recursively performs torch.chunk().

Parameters
  • input – Any input with tensors, tuples, lists, dicts, and other objects.

  • chunks (int) – Number of chunks to return.

  • dim (int) – Dimension along which to split the tensors. Defaults to 0.

Returns

List length chunks with sliced tensors.

Return type

list of Any

Example

>>> x = [torch.ones(4, 2),
...      {'target': torch.zeros(4), 'weights': torch.ones(4)}]
>>> x
[tensor([[1., 1.],
         [1., 1.],
         [1., 1.],
         [1., 1.]]),
 {'target': tensor([0., 0., 0., 0.]), 'weights': tensor([1., 1., 1., 1.])}]
>>> deep_chunk(x, 2, 0)
[[tensor([[1., 1.],
          [1., 1.]]),
  {'target': tensor([0., 0.]), 'weights': tensor([1., 1.])}],
 [tensor([[1., 1.],
          [1., 1.]]),
  {'target': tensor([0., 0.]), 'weights': tensor([1., 1.])}]]