torch.jit.freeze¶
- torch.jit.freeze(mod, preserved_attrs=None, optimize_numerics=True)[源代码]¶
冻结 ScriptModule,将子模块和属性内联为常量。
冻结
ScriptModule
将克隆它,并尝试将克隆模块的子模块、参数和属性内联为 TorchScript IR 图中的常量。默认情况下,将保留 forward 以及在 preserved_attrs 中指定的属性和方法。此外,在保留方法中修改的任何属性也将被保留。冻结目前只接受处于评估模式下的 ScriptModule。
冻结应用通用优化,无论机器如何,都会加快模型速度。要使用服务器特定的设置进行进一步优化,请在冻结后运行 optimize_for_inference。
- 参数
mod (
ScriptModule
) – 要冻结的模块preserved_attrs (Optional[List[str]]) – 除了前向方法之外,要保留的属性列表。在保留方法中修改的属性也将被保留。
optimize_numerics (bool) – 如果为
True
,将运行一组优化步骤,这些步骤不严格保留数值。优化详情请参阅 torch.jit.run_frozen_optimizations。
- 返回值
已冻结的
ScriptModule
。
示例(冻结具有参数的简单模块)
def forward(self, input): output = self.weight.mm(input) output = self.linear(output) return output scripted_module = torch.jit.script(MyModule(2, 3).eval()) frozen_module = torch.jit.freeze(scripted_module) # parameters have been removed and inlined into the Graph as constants assert len(list(frozen_module.named_parameters())) == 0 # See the compiled graph as Python code print(frozen_module.code)
示例(冻结具有保留属性的模块)
def forward(self, input): self.modified_tensor += 1 return input + self.modified_tensor scripted_module = torch.jit.script(MyModule2().eval()) frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["version"]) # we've manually preserved `version`, so it still exists on the frozen module and can be modified assert frozen_module.version == 1 frozen_module.version = 2 # `modified_tensor` is detected as being mutated in the forward, so freezing preserves # it to retain model semantics assert frozen_module(torch.tensor(1)) == torch.tensor(12) # now that we've run it once, the next result will be incremented by one assert frozen_module(torch.tensor(1)) == torch.tensor(13)
注意
还支持冻结子模块属性:frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=[“submodule.version”])
注意
如果您不确定为什么某个属性没有内联为常量,可以在 frozen_module.forward.graph 上运行 dump_alias_db,以查看冻结是否检测到该属性是否被修改。
注意
由于冻结会使权重成为常量并删除模块层次结构,因此 to 和其他用于操作设备或数据类型的 nn.Module 方法不再起作用。作为一种解决方法,您可以在 torch.jit.load 中指定 map_location 来重新映射设备,但特定于设备的逻辑可能已被烘焙到模型中。