This snippet be used to print a short overview of a complex data structure containing dicts, sequences and tensors

explain_structure.py

import torch

def is_iterable(obj):
  try:
    iter(obj)
  except TypeError:
    return False
  else:
    return not isinstance(obj, str)

def explain_structure(obj, all=False):
    # print(type(obj))
    if hasattr(obj, 'shape'):
        repr = f"({type(obj)}, shape={obj.shape})"
        if type(obj) == torch.Tensor:
            repr = f"{repr} {obj.device}"
        return repr
    elif type(obj) in [type(dict())]:
        return "{ " + ",".join([f"{key}: " + explain_structure(obj[key]) for key in obj.keys()]) + " }"
    elif type(obj) in [type([]), type(())] or is_iterable(obj):
        inner = obj
        if not all:
            inner = inner[:3]
        inner_explain = ", ".join([f"[{i}]= ( {explain_structure(inner_obj)} )" for i, inner_obj in enumerate(inner)])
        return f"{type(obj)}, len={len(obj)}, {inner_explain}"
    elif isinstance(obj, str) or isinstance(obj, float) or isinstance(obj, int):
        return str(obj)
    else:
        return f"{type(obj)}"