Visualizer-简化你的ViT可视化之源码分析

 Visualizer！简化你的Vision Transformer可视化！ - 知乎 (zhihu.com)

今天使用了这个库，很好用，使用方法原作者写了，可以查阅上面链接，强烈推荐

使用顺手了自然想看看源码，方便自己个性化修改，本以为是注册的hook实现的，但是发现代码不太对劲，用的是装饰器加反编译实现的，下面详细看一下

from bytecode import Bytecode, Instr

class get_local(object):
    cache = {}
    is_activate = False

    def __init__(self, varname):
        self.varname = varname

    def __call__(self, func):
        if not type(self).is_activate:
            return func

        type(self).cache[func.__qualname__] = []
        c = Bytecode.from_code(func.__code__)
        extra_code = [
                         Instr('STORE_FAST', '_res'),
                         Instr('LOAD_FAST', self.varname),
                         Instr('STORE_FAST', '_value'),
                         Instr('LOAD_FAST', '_res'),
                         Instr('LOAD_FAST', '_value'),
                         Instr('BUILD_TUPLE', 2),
                         Instr('STORE_FAST', '_result_tuple'),
                         Instr('LOAD_FAST', '_result_tuple'),
                     ]
        c[-1:-1] = extra_code
        func.__code__ = c.to_code()

        def wrapper(*args, **kwargs):
            res, values = func(*args, **kwargs)
            type(self).cache[func.__qualname__].append(values.detach().cpu().numpy())
            return res
        return wrapper

主要实现的逻辑就是上面代码，是不是很简洁

看看怎么做到的

定义类内存储

class get_local(object):

    cache = {}

    is_activate = False

首先在类里面定义了一个状态参数和一个空字典用于存储，

注意这个是在实例化之前就有的，是在类中定义的，所有实例化都可以访问

之后每个不同的实例化是为了捕捉不同的中间参数，统一放进类的字典cache中

初始实例化函数

    def __init__(self, varname):

        self.varname = varname

 实例化，varname是你想捕捉的变量名，比如ViT中注意力头的中间变量，atten

__call__方法

    def __call__(self, func):

 接下来看作者重写的__call__这个魔法函数

在Python中，函数其实是一个对象，特殊之处是他是可调用对象，内部实现了特殊方法__call__(）

一个类实例也可以变成一个可调用对象，只需要实现__call__

那这个class的实例化对象就可以作为函数使用，用在下面作为装饰器

作为装饰器

在作者说明的使用方法中，全文链接见文章开头

在模型文件里，我们这么写

from visualizer import get_local

@get_local('attention_map') # 我要拿attention_map这个变量,所以把他传参给get_local
def your_attention_function(*args, **kwargs):
    ...
    attention_map = ... 
    ...
    return ...

装饰器简单概念这里不展开了，可以查看“：python 装饰器详解 - 知乎 (zhihu.com)

下面代码主体结构如下，作用是

1. 通过 @logging(level="TEST") 的形式给函数添加装饰器，表示该函数需要被记录日志。

2. 当调用带 @logging 装饰器的函数时，会输出打印日志的内容："[{0}]: enter {1}()"。其中 self.level 表示在初始化日志对象时传入的参数 level，表示当前日志级别；func.__name__ 表示被装饰函数的名字。

3. 在打完日志后，会继续执行原本需要执行的函数，并返回执行结果。

class logging(object):
    def __init__(self, level):
        self.level = level   # 初始化日志级别

    def __call__(self, func):  # 实现装饰器需要使用 __call__
        def wrapper(*args, **kwargs):  # 定义包装器来接收函数和指定参数
            print("[{0}]: enter {1}()".format(self.level, func.__name__))  # 打印日志信息
            return func(*args, **kwargs)  # 继续执行函数，并返回结果
        return wrapper

@logging(level="TEST")
def hello(a, b, c):
    print(a, b, c)

hello("hello,","good","morning")

-----------------------------
结果：
>>>[TEST]: enter hello()
>>>hello, good morning

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            res, values = func(*args, **kwargs)
            type(self).cache[func.__qualname__].append(values.detach().cpu().numpy())
            return res
        return wrapper

visualizer的代码主体结构和上面一样，就是正常函数操作之前把values放进之前的cache字典里

问题是values是怎么来的呢？这里作者没有用hook，而是用的编译，很神奇

反编译

        type(self).cache[func.__qualname__] = []

        c = Bytecode.from_code(func.__code__)

        extra_code = [

                         Instr('STORE_FAST', '_res'),

                         Instr('LOAD_FAST', self.varname),

                         Instr('STORE_FAST', '_value'),

                         Instr('LOAD_FAST', '_res'),

                         Instr('LOAD_FAST', '_value'),

                         Instr('BUILD_TUPLE', 2),

                         Instr('STORE_FAST', '_result_tuple'),

                         Instr('LOAD_FAST', '_result_tuple'),

                     ]

        c[-1:-1] = extra_code

        func.__code__ = c.to_code()

 先贴个链接，开完组会再细讲

(26条消息) Python函数属性、__code__属性的解释、PyCodeObject_lamehd的博客-CSDN博客

python反编译之字节码 - 掘金 (juejin.cn)