把上一个阶段的工作收尾之后，终于有时间实现一些平时想实现但比较麻烦就一直懒得弄的功能。事情的起因是这样的：刚刚我正在研究 Human3.6M 数据集的标注，这个数据集的标注是一个类似 JSON 的格式。并且除了一般的 JSON 文件结构外，其中还有一些字段的值是高维 numpy 数组。我试着将其中的两项打印出来，命令行瞬间被一大堆各种各样的数值填满了。尽管 IPython Terminal 的输出有自带的格式化功能，但整个输出的格式还是被巨大的数组打乱了。

实际上我平时经常行地遇到这种问题，一般的解决方法是仿照 mmcv.parallel.collate 中的形式写一个递归打印函数，然后把诸如 np.ndarray、torch.Tensor 以及 mmcv.DataContainer 这种高维数据的 __repr__ 函数映射为一个形式类似 lambda x: f"{x.shape}, {x.dtype}" 的函数。

这种方式虽然能在一定程度上解决高维数据打印的问题，但是由于函数逻辑未经仔细考虑，打印的效果依然不是非常好，甚至在一些 corner case 处会出现十分诡异的效果。为了比较彻底地解决这个问题，我决定基于比较成熟的 Python built-in pprint 来拓展开发，实现这一功能。

我首先考虑的设计方式是引入类似 json.JSONEncoder 的设计方式，也就是在使用 pprint 函数时，添加一个额外的参数，将一个 formatter 类传入，对特定类型对象的序列化方式定义在这个类中，类似这样：

1
2
3
4
5
6
7
class CustomFormatter(DefaultFormatter):
    def default(self, obj):
        if isinstance(obj, (np.ndarray, torch.Tensor)):
            return f"<{obj.__class__.__name__} {obj.dtype} of shape {obj.shape}>"
        return DefaultFormatter.default(self, obj)

pprint.pprint(obj, cls=CustomFormatter)

然而观察了一下 pprint 的内部实现之后，我发现要改成这种形式还有一点难度。主要是因为在 pprint 内部，不同对象的序列化并非依赖于一个 formatter，而是通过维护一个字典，并将不同类型对象的序列化方法注册到这个字典中，在遇到相应对象时直接从字典中找到相应方法进行调用即可，类似这种形式：

1
2
3
4
5
6
7
8
9
_dispatch = {}

def _pprint_list(self, object, stream, indent, allowance, context, level):
    stream.write('[')
    self._format_items(object, stream, indent, allowance + 1,
                       context, level)
    stream.write(']')

_dispatch[list.__repr__] = _pprint_list

如此看来，要想将其拓展到更多类别，只需要把 np.ndarray、torch.Tensor 等类型的 __repr__ 函数也注册到这一 _dispatch 字典中即可。然而此时仍然存在一个问题，这个字典是被所有的 PrettyPrinter 对象共享的，因此如果直接在其中添加键值对，会影响 Python 内置方法的行为。我希望能够在维持 built-in 行为不变的前提下添加这一功能，因此在实际进行实现时，我采用了二级 dispatch 字典，在遇到过长对象时，首先在 _dispatch_override 中查询，然后再在原有的字典中查询，即可实现重载功能。

为了使这一类型具有较好的可拓展性，我也计划开放一个 register 接口，便于用户将更多自定义的类型注册到字典中，以提高更多场景下的可用性。

最后放一个输出的示例（目前还没有适配 mmcv.parallel.DataContainer，仅展示 ndarray 和 Tensor 的效果）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
{'address': {'city': 'New York',
             'postalCode': '10021',
             'state': 'NY',
             'streetAddress': '21 2nd Street'},
 'age': 25,
 'firstName': 'John',
 'lastName': 'Smith',
 'nestedDict': {'nestedList': [<Tensor (shape=torch.Size([233]) dtype=torch.float32)>,
                               <Tensor (shape=torch.Size([114]) dtype=torch.float32)>,
                               <ndarray (shape=(514,) dtype=float64)>],
                'nestedNumpyArray': <ndarray (shape=(123, 456) dtype=float64)>},
 'numpyArray': <ndarray (shape=(3, 4, 5, 6) dtype=float64)>,
 'phoneNumber': [{'number': '212 555-1234', 'type': 'home'},
                 {'number': '646 555-4567', 'type': 'fax'}],
 'sex': 'male',
 'torchTensor': <Tensor (shape=torch.Size([6, 7, 8]) dtype=torch.float32)>}

这一模块目前还在开发和测试阶段，预计会在测试稳定后整合进我的 pypi package 中发布，如果想使用预览版，可以参考这一链接。

2023-07-03 更新：前段时间发现了一种更好的实现方式，即基于 prettyprinter 实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import prettyprinter
import torch
import numpy as np

@prettyprinter.register_pretty(torch.Tensor)
def pretty_torchtensor(tensor: torch.Tensor, ctx):
    return prettyprinter.pretty_call(ctx,
                                     torch.Tensor,
                                     shape=tensor.shape,
                                     dtype=tensor.dtype,
                                     device=tensor.device)

@prettyprinter.register_pretty(np.ndarray)
def pretty_npndarray(ndarray: np.ndarray, ctx):
    return prettyprinter.pretty_call(ctx,
                                     np.ndarray,
                                     shape=ndarray.shape,
                                     dtype=ndarray.dtype)


pprint = prettyprinter.pprint

prettyprinter 的输出示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
{
    'numpyArray': numpy.ndarray(
        shape=(224, 224, 3),
        dtype=dtype('float64')
    ),
    'torchTensor': torch.Tensor(
        shape=torch.Size((64, 3, 224, 224)),
        dtype=torch.float32,
        device=device(type='cpu')
    ),
    'nestedDict': {
        'nestedArray': numpy.ndarray(
            shape=(224, 224, 3),
            dtype=dtype('float64')
        ),
        'nestedTensor': torch.Tensor(
            shape=torch.Size((64, 3, 224, 224)),
            dtype=torch.float32,
            device=device(type='cpu')
        )
    }
}