File: model_compatibility.h

package info (click to toggle)
pytorch-cuda 2.6.0%2Bdfsg-7
  • links: PTS, VCS
  • area: contrib
  • in suites: forky, sid, trixie
  • size: 161,620 kB
  • sloc: python: 1,278,832; cpp: 900,322; ansic: 82,710; asm: 7,754; java: 3,363; sh: 2,811; javascript: 2,443; makefile: 597; ruby: 195; xml: 84; objc: 68
file content (104 lines) | stat: -rw-r--r-- 3,612 bytes parent folder | download | duplicates (3)
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
#pragma once

#include <ATen/core/ivalue.h>
#include <c10/macros/Export.h>
#include <torch/csrc/jit/mobile/compatibility/runtime_compatibility.h>

#include <istream>
#include <memory>
#include <unordered_map>
#include <vector>

namespace caffe2::serialize {
class PyTorchStreamReader;
class ReadAdapterInterface;
} // namespace caffe2::serialize

namespace torch::jit {

// The family of methods below to get bytecode version from a model
// Throws if not passed in a well formed model
TORCH_API uint64_t _get_model_bytecode_version(std::istream& in);

TORCH_API uint64_t _get_model_bytecode_version(const std::string& filename);

TORCH_API uint64_t _get_model_bytecode_version(
    const std::shared_ptr<caffe2::serialize::ReadAdapterInterface>& rai);

uint64_t _get_model_bytecode_version(
    const std::vector<c10::IValue>& bytecode_ivalues);

// The family of methods below to get the operator version from a model
// Throws if not passed in a well formed model
TORCH_API uint64_t _get_model_operator_version(std::istream& in);

TORCH_API uint64_t _get_model_operator_version(const std::string& filename);

TORCH_API uint64_t _get_model_operator_version(
    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);

// Utility Functions
std::vector<c10::IValue> get_bytecode_ivalues(
    caffe2::serialize::PyTorchStreamReader& reader);

c10::IValue readArchive(
    const std::string& archive_name,
    caffe2::serialize::PyTorchStreamReader& stream_reader);

bool check_zip_file(
    const std::shared_ptr<caffe2::serialize::ReadAdapterInterface>& rai);

// The family of methods below to get the root ops and information from a model
TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
    std::istream& in);

TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
    const std::string& filename);

TORCH_API std::unordered_map<std::string, OperatorInfo> _get_model_ops_and_info(
    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);

// The family of methods below to get contained types from a model
// Throws if not passed in a well formed model
TORCH_API std::unordered_set<std::string> _get_mobile_model_contained_types(
    std::istream& in);

TORCH_API std::unordered_set<std::string> _get_mobile_model_contained_types(
    const std::string& filename);

TORCH_API std::unordered_set<std::string> _get_mobile_model_contained_types(
    std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);

std::unordered_set<std::string> _get_mobile_model_contained_types(
    const std::vector<c10::IValue>& bytecode_ivalues);

// The family of methods below return the compatibility information of a model
struct ModelCompatibilityInfo {
  uint64_t bytecode_version;
  std::unordered_map<std::string, OperatorInfo> operator_info;
  std::unordered_set<std::string> type_table;
  uint64_t operator_version;

  // Factory Methods
  static TORCH_API ModelCompatibilityInfo get(std::istream& in);
  static TORCH_API ModelCompatibilityInfo get(const std::string& filename);
  static TORCH_API ModelCompatibilityInfo
  get(std::shared_ptr<caffe2::serialize::ReadAdapterInterface> rai);
};

enum ModelCompatibilityStatus {
  OK = 1,
  ERROR = 2,
};

struct ModelCompatCheckResult {
  ModelCompatibilityStatus status;
  std::vector<std::string> errors{};
};
// Takes in information about a runtime and a model and returns if the two are
// compatible with one another.
TORCH_API ModelCompatCheckResult is_compatible(
    RuntimeCompatibilityInfo runtime_info,
    const ModelCompatibilityInfo& model_info);

} // namespace torch::jit