misc.py 文件源码

def shapes_all(data):
    """
    Recursively walks the data (can be tuples, lists, or dict) and
    replaces a tensor with its shape tuple whenever it meets a tensor
    """
    if isinstance(data, (tuple, list)):
        ans = map(shapes_all, data)
        return type(data)(ans)
    elif isinstance(data, dict):
        return {k: shapes_all(v) for k, v in data.items()}
    elif (isinstance(data, np.ndarray)
          or torch.is_tensor(data)
          or isinstance(data, torch.autograd.Variable)
          or isinstance(data, torch.nn.Parameter)):
        return shape(data)
    else:
        return data