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)
下面代码主体结构如下,作用是
-
通过
@logging(level="TEST")
的形式给函数添加装饰器,表示该函数需要被记录日志。 -
当调用带
@logging
装饰器的函数时,会输出打印日志的内容:"[{0}]: enter {1}()"
。其中self.level
表示在初始化日志对象时传入的参数level
,表示当前日志级别;func.__name__
表示被装饰函数的名字。 -
在打完日志后,会继续执行原本需要执行的函数,并返回执行结果。
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博客