在 VS Code 中调试 Tensor 形状不显示的问题及解决方案

文章目录

常见问题
解决方案
- 1. 定制类包装和 `__repr__` 方法
解释
如何应用
总结

在使用 VS Code 调试 PyTorch 代码时，可能会遇到一个常见问题：调试时 variables 窗口中不显示 Tensor 的形状信息。这会使得调试时观察数据的结构变得不便，尤其是在处理高维度的张量时。

在这篇博客中，我将提供一个改进方案，通过自定义 __repr__ 方法，让 PyTorch Tensor 的形状在调试信息中更清晰地呈现。我们还将展示如何针对其他 Python 对象（如 list、dict 以及 pandas.DataFrame）进行类似的处理，使调试时可以直观地看到这些对象的尺寸或内容。

常见问题

在调试 PyTorch 项目时，VS Code 的 variables 窗口默认不会显示张量的形状信息，除非你展开查看具体的变量内容。这种行为对于需要快速检查张量维度的场景是比较不便的。Python 的内置类型（如 bool、int、str）也无法直接继承，因此我们不能简单地扩展这些类型来自定义调试信息。

为了解决这些问题，我们可以通过创建包装类（wrapper classes）或直接重写对象的 __repr__ 方法，来定制它们在调试窗口中的显示内容。

解决方案

我们可以通过以下步骤来自定义调试时 Tensor 以及其他对象的显示方式。

1. 定制类包装和 `repr` 方法

我们可以为 Python 的内置类型创建包装类，重写 __repr__ 方法以自定义它们的显示形式。
对于 PyTorch 的 Tensor 和 pandas.DataFrame，可以直接重写它们的 __repr__ 方法来显示额外的信息（如张量形状或 DataFrame 的维度）。

下面是一个完整的代码示例：

import torch
import pandas as pd# -------------------- 自定义包装类 --------------------class CustomBool:def __init__(self, value):self.value = bool(value)def __repr__(self):return f'{{bool}} {self.value}'class CustomInt:def __init__(self, value):self.value = int(value)def __repr__(self):return f'{{int}} {self.value}'class CustomStr:def __init__(self, value):self.value = str(value)def __repr__(self):return f'{{str}} {self.value}'# 自定义 list 和 dict 子类
class CustomList(list):def __repr__(self):return f'{{list: {len(self)}}} {super().__repr__()}'class CustomDict(dict):def __repr__(self):return f'{{dict: {len(self)}}} {super().__repr__()}'# 自定义 Tensor 的 __repr__ (Torch)
original_tensor_repr = torch.Tensor.__repr__
def custom_tensor_repr(self):return f'{{Tensor: {tuple(self.shape)}}} {original_tensor_repr(self)}'
torch.Tensor.__repr__ = custom_tensor_repr# 自定义 DataFrame 的 __repr__ (Pandas)
original_dataframe_repr = pd.DataFrame.__repr__
def custom_dataframe_repr(self):return f'{{DataFrame: {self.shape}}} {original_dataframe_repr(self)}'
pd.DataFrame.__repr__ = custom_dataframe_repr# 自定义 DataLoader 的类
class DataLoader:def __init__(self, data_size):self.data_size = data_sizedef __len__(self):return self.data_sizedef __repr__(self):return f'{{DataLoader: {len(self)}}} DataLoader object'# -------------------- __main__ 函数 --------------------
def main():# 使用自定义类型代替原生类型my_list = CustomList([1, 2, 3, 4, 5, 6])my_dict = CustomDict({'a': 1, 'b': 2, 'c': 3})my_bool = CustomBool(True)my_int = CustomInt(42)my_str = CustomStr("hello")# 测试 Tensormy_tensor = torch.randn(100, 512)# 测试 DataFramemy_dataframe = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]})# 测试 DataLoadermy_dataloader = DataLoader(220)# 输出内容print(my_list)        # {list: 6} [1, 2, 3, 4, 5, 6]print(my_dict)        # {dict: 3} {'a': 1, 'b': 2, 'c': 3}print(my_bool)        # {bool} Trueprint(my_int)         # {int} 42print(my_str)         # {str} 'hello'print(my_tensor)      # {Tensor: (100, 512)} tensor([...])print(my_dataframe)   # {DataFrame: (3, 3)}    A  B  Cprint(my_dataloader)  # {DataLoader: 220} DataLoader object# 如果是直接运行文件，则调用 main 函数
if __name__ == "__main__":main()

解释

包装类的实现：
- 我们为 bool、int 和 str 创建了包装类 CustomBool、CustomInt 和 CustomStr，并通过重写 __repr__ 方法来定制它们的显示格式。
自定义 list 和 dict：
- 使用 CustomList 和 CustomDict 继承自 Python 的内置 list 和 dict 类型，并且通过重写 __repr__ 方法来显示列表或字典的长度。这对于调试包含大量数据的容器类非常有用。
Tensor 和 DataFrame 的自定义 __repr__：
- 我们通过直接修改 torch.Tensor 和 pandas.DataFrame 的 __repr__ 方法，分别让它们在打印时显示张量的形状和 DataFrame 的维度信息。这在调试过程中可以让你一眼就看到数据的结构，而不需要展开变量。
DataLoader 自定义类：
- 我们还创建了一个自定义的 DataLoader 类，并重写了它的 __repr__ 方法来显示加载的数据大小。这对于在训练模型时检查数据集大小十分有用。