1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
/** \brief This file defines freezing Torchscript module API.
*
* This API has python-binding and can be invoked directly or as a part of
* general optimization pipeline.
*/
#pragma once
#include <torch/csrc/jit/api/module.h>
#include <torch/csrc/jit/ir/ir.h>
/** \brief Freeze Module, i.e., Assume all atrributes are constants.
*
* Freezing module is a functionality that allows the JIT to internalize
* immutable attributes. Combined with inlining, the module is aggressively
* optimized and significant overhead is optimized away. The freezeModule API
* produces a cloned frozen module.
*/
namespace torch {
namespace jit {
TORCH_API Module freeze_module(
const Module& module,
std::vector<std::string> preservedAttrs = std::vector<std::string>(),
bool freezeInterfaces = true,
bool preserveParameters = false);
// Clone-free version of freeze_module. This modifies the module inplace.
// Use this version to avoid extra memory usage incurred by cloning the module.
TORCH_API void freeze_module_inplace(
Module* module,
std::vector<std::string> preservedAttrs = std::vector<std::string>(),
bool freezeInterfaces = true,
bool preserveParameters = false);
} // namespace jit
} // namespace torch
|
