Documentation
¶
Index ¶
- func ExportONNX(model nn.Module, sampleInput *tensor.Tensor, filePath string) error
- func ExportTorchScript(model nn.Module, sampleInput *tensor.Tensor, filePath string) error
- func SaveTorchScript(ts *TorchScriptModel, filePath string) error
- type ONNXGraph
- type ONNXModel
- type ONNXNode
- type ONNXTensorInfo
- type TSNode
- type TorchScriptModel
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func ExportONNX ¶
ExportONNX traces the module with sampleInput and writes an ONNX graph file to disk.
func ExportTorchScript ¶
ExportTorchScript traces a model and saves it as a TorchScript-compatible package in one call.
func SaveTorchScript ¶
func SaveTorchScript(ts *TorchScriptModel, filePath string) error
SaveTorchScript writes the TorchScript model into a standard PKZip package (.pt format).
Types ¶
type ONNXGraph ¶
type ONNXGraph struct {
Name string `json:"name"`
Nodes []ONNXNode `json:"nodes"`
Inputs []ONNXTensorInfo `json:"inputs"`
Outputs []ONNXTensorInfo `json:"outputs"`
Initializers []ONNXTensorInfo `json:"initializers"`
}
ONNXGraph holds nodes, inputs, outputs, and weight initializers.
type ONNXModel ¶
type ONNXModel struct {
IRVersion int `json:"ir_version"`
ProducerName string `json:"producer_name"`
Graph ONNXGraph `json:"graph"`
}
ONNXModel encapsulates an exported computational graph.
type ONNXNode ¶
type ONNXNode struct {
OpType string `json:"op_type"`
Inputs []string `json:"inputs"`
Outputs []string `json:"outputs"`
Attributes map[string]interface{} `json:"attributes,omitempty"`
}
ONNXNode represents an operational node in an ONNX computation graph.
type ONNXTensorInfo ¶
type ONNXTensorInfo struct {
Name string `json:"name"`
Shape []int `json:"shape"`
Data []float32 `json:"data,omitempty"`
}
ONNXTensorInfo records shape and type for inputs, outputs, and initializers.
type TSNode ¶
type TSNode struct {
OpType string `json:"op"`
Inputs []string `json:"inputs"`
Outputs []string `json:"outputs"`
}
TSNode represents a traced operation node in a TorchScript computation graph.
type TorchScriptModel ¶
type TorchScriptModel struct {
Name string `json:"name"`
Inputs []string `json:"inputs"`
Outputs []string `json:"outputs"`
Nodes []TSNode `json:"nodes"`
Parameters map[string]*tensor.Tensor `json:"-"`
}
TorchScriptModel holds the traced graph and weights for export.
func TraceTorchScript ¶
TraceTorchScript executes a forward pass and traces the operation flow.
Source Files