序列化前修改对象:类C父子进程GPU/CPU切换的递归问题
解决类C序列化时GPU/CPU分离的无限递归问题
首先得搞清楚你碰到的无限递归为啥会发生:当你在__reduce__里调用self.copy(shallow=True)的时候,如果这个copy方法底层用到了pickle(或者其他依赖对象序列化的机制),那它又会触发__reduce__方法,直接陷入“调用copy→触发__reduce__→再调用copy”的死循环里了。
下面是两种可行的解决思路,你可以根据自己的代码结构选:
思路一:手动浅复制属性,避免触发序列化
既然copy方法会触发递归,那我们就不用它,直接手动复制对象的属性,完全绕开序列化流程:
def __reduce__(self): """避免子进程使用GPU的序列化处理""" # 手动创建一个类C的空实例,不触发__init__里的GPU初始化 p = object.__new__(C) # 浅复制所有实例属性 p.__dict__.update(self.__dict__) # 设置CPU-only标志 p._cpu_only = True return _dummy_rebuild, (p,)
这里用object.__new__(C)创建空实例,不会触发类的__init__方法(避免不必要的GPU资源初始化),然后直接复制__dict__里的属性,完全不会触发序列化相关逻辑,自然就不会递归调用__reduce__了。
思路二:临时禁用__reduce__方法再执行copy
如果你的copy方法有其他必须保留的逻辑,不能手动复制属性,那可以在执行copy前临时把__reduce__方法替换掉,完成复制后再恢复:
def __reduce__(self): """避免子进程使用GPU的序列化处理""" # 临时保存原__reduce__方法,替换成默认的序列化逻辑 original_reduce = self.__class__.__reduce__ self.__class__.__reduce__ = object.__reduce__ try: # 现在执行copy不会触发自定义的__reduce__了 p = self.copy(shallow=True) p._cpu_only = True finally: # 不管有没有异常,都要恢复原__reduce__方法 self.__class__.__reduce__ = original_reduce return _dummy_rebuild, (p,)
这个方法的核心是在copy过程中让类使用默认的object.__reduce__,避免触发我们自定义的逻辑,等copy完成后再把原方法恢复,保证后续的序列化行为正常。
另外要注意_dummy_rebuild函数的实现,它需要正确重建对象,比如:
def _dummy_rebuild(obj): return obj
这样就能确保子进程拿到的是已经设置了_cpu_only=True的实例,只会使用CPU资源了。
内容的提问来源于stack exchange,提问作者Kiuhnm
相关产品推荐
相关产品推荐

