statedict2pytree 0.5.0__py3-none-any.whl → 0.5.3__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- client/.gitignore +3 -0
- client/package-lock.json +4540 -0
- client/package.json +36 -0
- client/public/bundle.js +10072 -0
- client/public/bundle.js.map +1 -0
- client/public/index.html +14 -0
- {statedict2pytree/static → client/public}/output.css +184 -109
- client/rollup.config.mjs +44 -0
- client/src/App.svelte +583 -0
- client/src/empty.ts +0 -0
- client/src/main.js +8 -0
- client/tailwind.config.js +8 -0
- client/tsconfig.json +5 -0
- statedict2pytree-0.5.3.dist-info/METADATA +77 -0
- statedict2pytree-0.5.3.dist-info/RECORD +17 -0
- statedict2pytree/__init__.py +0 -7
- statedict2pytree/statedict2pytree.py +0 -219
- statedict2pytree/templates/index.html +0 -308
- statedict2pytree-0.5.0.dist-info/METADATA +0 -147
- statedict2pytree-0.5.0.dist-info/RECORD +0 -8
- {statedict2pytree/static → client/public}/input.css +0 -0
- {statedict2pytree-0.5.0.dist-info → statedict2pytree-0.5.3.dist-info}/WHEEL +0 -0
statedict2pytree/__init__.py
DELETED
|
@@ -1,219 +0,0 @@
|
|
|
1
|
-
import functools as ft
|
|
2
|
-
import re
|
|
3
|
-
|
|
4
|
-
import equinox as eqx
|
|
5
|
-
import flask
|
|
6
|
-
import jax
|
|
7
|
-
import numpy as np
|
|
8
|
-
from beartype.typing import Optional
|
|
9
|
-
from jaxtyping import PyTree
|
|
10
|
-
from loguru import logger
|
|
11
|
-
from penzai import pz
|
|
12
|
-
from pydantic import BaseModel
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
app = flask.Flask(__name__)
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
class Field(BaseModel):
|
|
19
|
-
path: str
|
|
20
|
-
shape: tuple[int, ...]
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
class TorchField(Field):
|
|
24
|
-
pass
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
class JaxField(Field):
|
|
28
|
-
type: str
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
PYTREE: Optional[PyTree] = None
|
|
32
|
-
STATE_DICT: Optional[dict] = None
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
def can_reshape(shape1, shape2):
|
|
36
|
-
product1 = np.prod(shape1)
|
|
37
|
-
product2 = np.prod(shape2)
|
|
38
|
-
|
|
39
|
-
return product1 == product2
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
def get_node(
|
|
43
|
-
tree: PyTree, targets: list[str], log_when_not_found: bool = False
|
|
44
|
-
) -> PyTree | None:
|
|
45
|
-
if len(targets) == 0 or tree is None:
|
|
46
|
-
return tree
|
|
47
|
-
else:
|
|
48
|
-
next_target: str = targets[0]
|
|
49
|
-
if bool(re.search(r"\[\d+\]", next_target)):
|
|
50
|
-
split_index = next_target.rfind("[")
|
|
51
|
-
name, index = next_target[:split_index], next_target[split_index:]
|
|
52
|
-
index = index[1:-1]
|
|
53
|
-
if hasattr(tree, name):
|
|
54
|
-
subtree = getattr(tree, name)[int(index)]
|
|
55
|
-
else:
|
|
56
|
-
subtree = None
|
|
57
|
-
if log_when_not_found:
|
|
58
|
-
logger.info(f"Couldn't find {name} in {tree.__class__}")
|
|
59
|
-
else:
|
|
60
|
-
if hasattr(tree, next_target):
|
|
61
|
-
subtree = getattr(tree, next_target)
|
|
62
|
-
else:
|
|
63
|
-
subtree = None
|
|
64
|
-
if log_when_not_found:
|
|
65
|
-
logger.info(f"Couldn't find {next_target} in {tree.__class__}")
|
|
66
|
-
return get_node(subtree, targets[1:])
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
def pytree_to_fields(pytree: PyTree) -> list[JaxField]:
|
|
70
|
-
flattened, _ = jax.tree_util.tree_flatten_with_path(pytree)
|
|
71
|
-
fields: list[JaxField] = []
|
|
72
|
-
for key_path, value in flattened:
|
|
73
|
-
path = jax.tree_util.keystr(key_path)
|
|
74
|
-
type_path = path.split(".")[1:-1]
|
|
75
|
-
target_path = path.split(".")[1:]
|
|
76
|
-
node_type = type(get_node(pytree, type_path, log_when_not_found=True))
|
|
77
|
-
node = get_node(pytree, target_path, log_when_not_found=True)
|
|
78
|
-
if node is not None and hasattr(node, "shape") and len(node.shape) > 0:
|
|
79
|
-
fields.append(
|
|
80
|
-
JaxField(path=path, type=str(node_type), shape=tuple(node.shape))
|
|
81
|
-
)
|
|
82
|
-
|
|
83
|
-
return fields
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
def state_dict_to_fields(state_dict: Optional[dict]) -> list[TorchField]:
|
|
87
|
-
if state_dict is None:
|
|
88
|
-
return []
|
|
89
|
-
fields: list[TorchField] = []
|
|
90
|
-
for key, value in state_dict.items():
|
|
91
|
-
if hasattr(value, "shape") and len(value.shape) > 0:
|
|
92
|
-
fields.append(TorchField(path=key, shape=tuple(value.shape)))
|
|
93
|
-
return fields
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
@app.route("/visualize", methods=["POST"])
|
|
97
|
-
def visualize_with_penzai():
|
|
98
|
-
global PYTREE, STATE_DICT
|
|
99
|
-
if PYTREE is None or STATE_DICT is None:
|
|
100
|
-
return flask.jsonify({"error": "No Pytree or StateDict found"})
|
|
101
|
-
request_data = flask.request.json
|
|
102
|
-
if request_data is None:
|
|
103
|
-
return flask.jsonify({"error": "No data received"})
|
|
104
|
-
jax_fields = request_data["jaxFields"]
|
|
105
|
-
torch_fields = request_data["torchFields"]
|
|
106
|
-
model, state = convert(jax_fields, torch_fields, PYTREE, STATE_DICT)
|
|
107
|
-
with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
|
|
108
|
-
html_jax = pz.ts.render_to_html((model, state))
|
|
109
|
-
html_torch = pz.ts.render_to_html(STATE_DICT)
|
|
110
|
-
|
|
111
|
-
combined_html = f"<html><body>{html_jax}<hr>{html_torch}</body></html>"
|
|
112
|
-
return combined_html
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
@app.route("/convert", methods=["POST"])
|
|
116
|
-
def convert_torch_to_jax():
|
|
117
|
-
global PYTREE, STATE_DICT
|
|
118
|
-
if PYTREE is None or STATE_DICT is None:
|
|
119
|
-
return flask.jsonify({"error": "No Pytree or StateDict found"})
|
|
120
|
-
request_data = flask.request.json
|
|
121
|
-
if request_data is None:
|
|
122
|
-
return flask.jsonify({"error": "No data received"})
|
|
123
|
-
|
|
124
|
-
jax_fields_json = request_data["jaxFields"]
|
|
125
|
-
jax_fields: list[JaxField] = []
|
|
126
|
-
for f in jax_fields_json:
|
|
127
|
-
shape_tuple = tuple(
|
|
128
|
-
[int(i) for i in f["shape"].strip("()").split(",") if len(i) > 0]
|
|
129
|
-
)
|
|
130
|
-
jax_fields.append(JaxField(path=f["path"], type=f["type"], shape=shape_tuple))
|
|
131
|
-
|
|
132
|
-
torch_fields_json = request_data["torchFields"]
|
|
133
|
-
torch_fields: list[TorchField] = []
|
|
134
|
-
for f in torch_fields_json:
|
|
135
|
-
shape_tuple = tuple(
|
|
136
|
-
[int(i) for i in f["shape"].strip("()").split(",") if len(i) > 0]
|
|
137
|
-
)
|
|
138
|
-
torch_fields.append(TorchField(path=f["path"], shape=shape_tuple))
|
|
139
|
-
|
|
140
|
-
name = request_data["name"]
|
|
141
|
-
model, state = convert(jax_fields, torch_fields, PYTREE, STATE_DICT)
|
|
142
|
-
eqx.tree_serialise_leaves(name, (model, state))
|
|
143
|
-
|
|
144
|
-
return flask.jsonify({"status": "success"})
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
@app.route("/", methods=["GET"])
|
|
148
|
-
def index():
|
|
149
|
-
pytree_fields = pytree_to_fields(PYTREE)
|
|
150
|
-
return flask.render_template(
|
|
151
|
-
"index.html",
|
|
152
|
-
pytree_fields=pytree_fields,
|
|
153
|
-
torch_fields=state_dict_to_fields(STATE_DICT),
|
|
154
|
-
)
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
def autoconvert(pytree: PyTree, state_dict: dict) -> tuple[PyTree, eqx.nn.State]:
|
|
158
|
-
jax_fields = pytree_to_fields(pytree)
|
|
159
|
-
torch_fields = state_dict_to_fields(state_dict)
|
|
160
|
-
|
|
161
|
-
for k, v in state_dict.items():
|
|
162
|
-
state_dict[k] = v.numpy()
|
|
163
|
-
return convert(jax_fields, torch_fields, pytree, state_dict)
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
def convert(
|
|
167
|
-
jax_fields: list[JaxField],
|
|
168
|
-
torch_fields: list[TorchField],
|
|
169
|
-
pytree: PyTree,
|
|
170
|
-
state_dict: dict,
|
|
171
|
-
) -> tuple[PyTree, eqx.nn.State]:
|
|
172
|
-
identity = lambda *args, **kwargs: pytree
|
|
173
|
-
model, state = eqx.nn.make_with_state(identity)()
|
|
174
|
-
state_paths: list[tuple[JaxField, TorchField]] = []
|
|
175
|
-
for jax_field, torch_field in zip(jax_fields, torch_fields):
|
|
176
|
-
if not can_reshape(jax_field.shape, torch_field.shape):
|
|
177
|
-
raise ValueError(
|
|
178
|
-
"Fields have incompatible shapes!"
|
|
179
|
-
f"{jax_field.shape=} != {torch_field.shape=}"
|
|
180
|
-
)
|
|
181
|
-
path = jax_field.path.split(".")[1:]
|
|
182
|
-
if "StateIndex" in jax_field.type:
|
|
183
|
-
state_paths.append((jax_field, torch_field))
|
|
184
|
-
|
|
185
|
-
else:
|
|
186
|
-
where = ft.partial(get_node, targets=path)
|
|
187
|
-
if where(model) is not None:
|
|
188
|
-
model = eqx.tree_at(
|
|
189
|
-
where,
|
|
190
|
-
model,
|
|
191
|
-
state_dict[torch_field.path].reshape(jax_field.shape),
|
|
192
|
-
)
|
|
193
|
-
result: dict[str, list[TorchField]] = {}
|
|
194
|
-
for tuple_item in state_paths:
|
|
195
|
-
path_prefix = tuple_item[0].path.split(".")[1:-1]
|
|
196
|
-
prefix_key = ".".join(path_prefix)
|
|
197
|
-
if prefix_key not in result:
|
|
198
|
-
result[prefix_key] = []
|
|
199
|
-
result[prefix_key].append(tuple_item[1])
|
|
200
|
-
|
|
201
|
-
for key in result:
|
|
202
|
-
state_index = get_node(model, key.split("."))
|
|
203
|
-
if state_index is not None:
|
|
204
|
-
to_replace_tuple = tuple([state_dict[i.path] for i in result[key]])
|
|
205
|
-
state = state.set(state_index, to_replace_tuple)
|
|
206
|
-
return model, state
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
def start_conversion(pytree: PyTree, state_dict: dict):
|
|
210
|
-
global PYTREE, STATE_DICT
|
|
211
|
-
if state_dict is None:
|
|
212
|
-
raise ValueError("STATE_DICT must not be None!")
|
|
213
|
-
PYTREE = pytree
|
|
214
|
-
STATE_DICT = state_dict
|
|
215
|
-
|
|
216
|
-
for k, v in STATE_DICT.items():
|
|
217
|
-
STATE_DICT[k] = v.numpy()
|
|
218
|
-
app.jinja_env.globals.update(enumerate=enumerate)
|
|
219
|
-
app.run(debug=True, port=5500)
|
|
@@ -1,308 +0,0 @@
|
|
|
1
|
-
<!doctype html>
|
|
2
|
-
<html lang="en" data-theme="light">
|
|
3
|
-
<head>
|
|
4
|
-
<meta charset="UTF-8" />
|
|
5
|
-
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
|
6
|
-
<meta http-equiv="X-UA-Compatible" content="ie=edge" />
|
|
7
|
-
<title>Torch2Jax</title>
|
|
8
|
-
|
|
9
|
-
<link
|
|
10
|
-
rel="stylesheet"
|
|
11
|
-
href="{{ url_for('static', filename='output.css') }}"
|
|
12
|
-
/>
|
|
13
|
-
<script src="https://cdn.jsdelivr.net/npm/sweetalert2@11"></script>
|
|
14
|
-
<script src="https://cdn.jsdelivr.net/npm/sortablejs@latest/Sortable.min.js"></script>
|
|
15
|
-
<script type="module">
|
|
16
|
-
let jaxSortable = new Sortable(document.getElementById("jax-fields"), {
|
|
17
|
-
animation: 150,
|
|
18
|
-
ghostClass: "blue-background-class",
|
|
19
|
-
});
|
|
20
|
-
|
|
21
|
-
let torchSortable = new Sortable(
|
|
22
|
-
document.getElementById("torch-fields"),
|
|
23
|
-
{
|
|
24
|
-
animation: 150,
|
|
25
|
-
ghostClass: "bg-blue-400",
|
|
26
|
-
onEnd: function (evt) {
|
|
27
|
-
document.getElementById("error-field").classList.add("hidden");
|
|
28
|
-
let allJaxFields = document.querySelectorAll("#jax-fields > div");
|
|
29
|
-
let allTorchFields = document.querySelectorAll(
|
|
30
|
-
"#torch-fields > div",
|
|
31
|
-
);
|
|
32
|
-
if (allJaxFields.length !== allTorchFields.length) {
|
|
33
|
-
Swal.fire({
|
|
34
|
-
icon: "error",
|
|
35
|
-
title:
|
|
36
|
-
"The number of fields in JAX and PyTorch should be the same",
|
|
37
|
-
});
|
|
38
|
-
} else {
|
|
39
|
-
for (let i = 0; i < allJaxFields.length; i++) {
|
|
40
|
-
let jaxField = allJaxFields[i];
|
|
41
|
-
let torchField = allTorchFields[i];
|
|
42
|
-
let jaxShape = jaxField.getAttribute("data-shape");
|
|
43
|
-
let torchShape = torchField.getAttribute("data-shape");
|
|
44
|
-
|
|
45
|
-
jaxShape = jaxShape
|
|
46
|
-
.replace("(", "")
|
|
47
|
-
.replace(")", "")
|
|
48
|
-
.replace(/\s+/g, "")
|
|
49
|
-
.replace(/,\s*$/, "");
|
|
50
|
-
|
|
51
|
-
torchShape = torchShape
|
|
52
|
-
.replace("(", "")
|
|
53
|
-
.replace(")", "")
|
|
54
|
-
.replace(/\s+/g, "")
|
|
55
|
-
.replace(/,\s*$/, "");
|
|
56
|
-
|
|
57
|
-
let jaxShapeParts = jaxShape.split(",").map((x) => parseInt(x));
|
|
58
|
-
let torchShapeParts = torchShape
|
|
59
|
-
.split(",")
|
|
60
|
-
.map((x) => parseInt(x));
|
|
61
|
-
let jaxShapeProduct = jaxShapeParts.reduce((a, b) => a * b, 1);
|
|
62
|
-
let torchShapeProduct = torchShapeParts.reduce(
|
|
63
|
-
(a, b) => a * b,
|
|
64
|
-
1,
|
|
65
|
-
);
|
|
66
|
-
|
|
67
|
-
let jaxEl = jaxField;
|
|
68
|
-
let torchEl = torchField;
|
|
69
|
-
if (jaxShapeProduct !== torchShapeProduct) {
|
|
70
|
-
jaxEl.classList.add("bg-error");
|
|
71
|
-
torchEl.classList.add("bg-error");
|
|
72
|
-
} else {
|
|
73
|
-
jaxEl.classList.remove("bg-error");
|
|
74
|
-
torchEl.classList.remove("bg-error");
|
|
75
|
-
}
|
|
76
|
-
}
|
|
77
|
-
}
|
|
78
|
-
},
|
|
79
|
-
},
|
|
80
|
-
);
|
|
81
|
-
</script>
|
|
82
|
-
</head>
|
|
83
|
-
<body class="w-10/12 mx-auto">
|
|
84
|
-
<script>
|
|
85
|
-
const Toast = Swal.mixin({
|
|
86
|
-
toast: true,
|
|
87
|
-
position: "top-end",
|
|
88
|
-
showConfirmButton: false,
|
|
89
|
-
timer: 5000,
|
|
90
|
-
timerProgressBar: true,
|
|
91
|
-
didOpen: (toast) => {
|
|
92
|
-
toast.onmouseenter = Swal.stopTimer;
|
|
93
|
-
toast.onmouseleave = Swal.resumeTimer;
|
|
94
|
-
},
|
|
95
|
-
});
|
|
96
|
-
|
|
97
|
-
async function visualize() {
|
|
98
|
-
const fields = getJaxAndTorchFields();
|
|
99
|
-
if (fields.error) {
|
|
100
|
-
Toast.fire({
|
|
101
|
-
icon: "error",
|
|
102
|
-
title: "The number of fields in JAX and PyTorch should be the same",
|
|
103
|
-
});
|
|
104
|
-
}
|
|
105
|
-
const jaxFields = fields.jaxFields;
|
|
106
|
-
const torchFields = fields.torchFields;
|
|
107
|
-
var data = JSON.stringify({
|
|
108
|
-
jaxFields: jaxFields,
|
|
109
|
-
torchFields: torchFields,
|
|
110
|
-
});
|
|
111
|
-
var xhr = new XMLHttpRequest();
|
|
112
|
-
xhr.open("POST", "/visualize", true);
|
|
113
|
-
xhr.setRequestHeader("Content-Type", "application/json");
|
|
114
|
-
xhr.onload = function () {
|
|
115
|
-
if (xhr.status >= 200 && xhr.status < 300) {
|
|
116
|
-
var container = document.getElementById("visualizationResult");
|
|
117
|
-
container.innerHTML = xhr.responseText;
|
|
118
|
-
var scripts = container.getElementsByTagName("script");
|
|
119
|
-
for (var i = 0; i < scripts.length; i++) {
|
|
120
|
-
var script = document.createElement("script");
|
|
121
|
-
script.text = scripts[i].text;
|
|
122
|
-
document.head.appendChild(script).parentNode.removeChild(script);
|
|
123
|
-
}
|
|
124
|
-
} else {
|
|
125
|
-
document.getElementById("visualizationResult").innerHTML =
|
|
126
|
-
"Error: " + xhr.statusText;
|
|
127
|
-
}
|
|
128
|
-
};
|
|
129
|
-
xhr.send(data);
|
|
130
|
-
}
|
|
131
|
-
|
|
132
|
-
function getJaxAndTorchFields() {
|
|
133
|
-
const jaxFields = Array.from(
|
|
134
|
-
document.querySelectorAll("#jax-fields")[0].children,
|
|
135
|
-
).map((li) => {
|
|
136
|
-
const path = li.getAttribute("data-path");
|
|
137
|
-
const shape = li.getAttribute("data-shape");
|
|
138
|
-
const type = li.getAttribute("data-type");
|
|
139
|
-
return { path, shape, type };
|
|
140
|
-
});
|
|
141
|
-
|
|
142
|
-
const torchFields = Array.from(
|
|
143
|
-
document.querySelectorAll("#torch-fields")[0].children,
|
|
144
|
-
).map((li) => {
|
|
145
|
-
const path = li.getAttribute("data-path");
|
|
146
|
-
const shape = li.getAttribute("data-shape");
|
|
147
|
-
return { path, shape };
|
|
148
|
-
});
|
|
149
|
-
|
|
150
|
-
const jaxLength = jaxFields.length;
|
|
151
|
-
const torchLength = torchFields.length;
|
|
152
|
-
if (jaxLength !== torchLength) {
|
|
153
|
-
return {
|
|
154
|
-
error: "The number of fields in JAX and PyTorch should be the same",
|
|
155
|
-
};
|
|
156
|
-
}
|
|
157
|
-
|
|
158
|
-
console.log({
|
|
159
|
-
jaxFields,
|
|
160
|
-
torchFields,
|
|
161
|
-
});
|
|
162
|
-
|
|
163
|
-
for (let i = 0; i < jaxLength; i++) {
|
|
164
|
-
if (jaxFields[i].shape !== torchFields[i].shape) {
|
|
165
|
-
Toast.fire({
|
|
166
|
-
icon: "error",
|
|
167
|
-
title: `${jaxFields[i].path} has shape ${jaxFields[i].shape}, while ${torchFields[i].path} has shape ${torchFields[i].shape}`,
|
|
168
|
-
});
|
|
169
|
-
document.getElementById("error-field").classList.remove("hidden");
|
|
170
|
-
document
|
|
171
|
-
.getElementById("error-field")
|
|
172
|
-
.querySelector("span").innerText =
|
|
173
|
-
`${jaxFields[i].path} has shape ${jaxFields[i].shape}, while ${torchFields[i].path} has shape ${torchFields[i].shape}`;
|
|
174
|
-
return { error: "Invalid shapes" };
|
|
175
|
-
}
|
|
176
|
-
}
|
|
177
|
-
|
|
178
|
-
return { jaxFields: jaxFields, torchFields: torchFields };
|
|
179
|
-
}
|
|
180
|
-
|
|
181
|
-
async function convert() {
|
|
182
|
-
const fields = getJaxAndTorchFields();
|
|
183
|
-
if (fields.error) {
|
|
184
|
-
Toast.fire({
|
|
185
|
-
icon: "error",
|
|
186
|
-
title: "Failed to convert!",
|
|
187
|
-
text: fields.error,
|
|
188
|
-
});
|
|
189
|
-
}
|
|
190
|
-
const jaxFields = fields.jaxFields;
|
|
191
|
-
const torchFields = fields.torchFields;
|
|
192
|
-
|
|
193
|
-
let idField = document.getElementById("name");
|
|
194
|
-
if (!idField) {
|
|
195
|
-
Toast.fire({
|
|
196
|
-
icon: "error",
|
|
197
|
-
title: "Error finding the name!",
|
|
198
|
-
});
|
|
199
|
-
}
|
|
200
|
-
|
|
201
|
-
let name = idField.value;
|
|
202
|
-
|
|
203
|
-
const response = await fetch("/convert", {
|
|
204
|
-
method: "POST",
|
|
205
|
-
headers: {
|
|
206
|
-
"Content-Type": "application/json",
|
|
207
|
-
},
|
|
208
|
-
body: JSON.stringify({
|
|
209
|
-
jaxFields,
|
|
210
|
-
torchFields,
|
|
211
|
-
name,
|
|
212
|
-
}),
|
|
213
|
-
});
|
|
214
|
-
|
|
215
|
-
const res = await response.json();
|
|
216
|
-
console.log(res);
|
|
217
|
-
if (res.error) {
|
|
218
|
-
Toast.fire({
|
|
219
|
-
icon: "error",
|
|
220
|
-
title: res.error,
|
|
221
|
-
});
|
|
222
|
-
} else {
|
|
223
|
-
Toast.fire({
|
|
224
|
-
icon: "success",
|
|
225
|
-
title: "Conversion successful",
|
|
226
|
-
});
|
|
227
|
-
}
|
|
228
|
-
}
|
|
229
|
-
</script>
|
|
230
|
-
<h1 class="text-3xl my-12">Welcome to Torch2Jax</h1>
|
|
231
|
-
|
|
232
|
-
<div class="grid grid-cols-2 gap-x-2">
|
|
233
|
-
<div class="">
|
|
234
|
-
<h2 class="text-2xl">JAX</h2>
|
|
235
|
-
<div id="jax-fields" class="bg-base-200">
|
|
236
|
-
{% for field in pytree_fields %}
|
|
237
|
-
<div
|
|
238
|
-
data-path="{{field.path}}"
|
|
239
|
-
data-shape="{{field.shape}}"
|
|
240
|
-
data-type="{{field.type}}"
|
|
241
|
-
class="whitespace-nowrap overflow-x-scroll cursor-pointer"
|
|
242
|
-
>
|
|
243
|
-
{{ field.path }} {{ field.shape }}
|
|
244
|
-
</div>
|
|
245
|
-
{% endfor %}
|
|
246
|
-
</div>
|
|
247
|
-
</div>
|
|
248
|
-
|
|
249
|
-
<div class="">
|
|
250
|
-
<h2 class="text-2xl">PyTorch</h2>
|
|
251
|
-
<div id="torch-fields" class="bg-base-200">
|
|
252
|
-
{% for field in torch_fields %}
|
|
253
|
-
<div
|
|
254
|
-
data-path="{{field.path}}"
|
|
255
|
-
data-shape="{{field.shape}}"
|
|
256
|
-
class="whitespace-nowrap overflow-x-scroll cursor-pointer"
|
|
257
|
-
>
|
|
258
|
-
{{ field.path }} {{ field.shape }}
|
|
259
|
-
</div>
|
|
260
|
-
{% endfor %}
|
|
261
|
-
</div>
|
|
262
|
-
</div>
|
|
263
|
-
</div>
|
|
264
|
-
<div class="flex justify-center my-12 w-full">
|
|
265
|
-
<div class="flex flex-col justify-center w-full">
|
|
266
|
-
<input
|
|
267
|
-
id="name"
|
|
268
|
-
type="text"
|
|
269
|
-
name="name"
|
|
270
|
-
class="input input-primary w-full"
|
|
271
|
-
placeholder="Name of the new file (model.eqx per default)"
|
|
272
|
-
value="model.eqx"
|
|
273
|
-
/>
|
|
274
|
-
<button
|
|
275
|
-
onclick="convert()"
|
|
276
|
-
class="btn btn-accent btn-wide btn-lg mx-auto my-2"
|
|
277
|
-
>
|
|
278
|
-
Convert!
|
|
279
|
-
</button>
|
|
280
|
-
</div>
|
|
281
|
-
</div>
|
|
282
|
-
<div role="alert" class="alert alert-error hidden" id="error-field">
|
|
283
|
-
<svg
|
|
284
|
-
xmlns="http://www.w3.org/2000/svg"
|
|
285
|
-
class="stroke-current shrink-0 h-6 w-6"
|
|
286
|
-
fill="none"
|
|
287
|
-
viewBox="0 0 24 24"
|
|
288
|
-
>
|
|
289
|
-
<path
|
|
290
|
-
stroke-linecap="round"
|
|
291
|
-
stroke-linejoin="round"
|
|
292
|
-
stroke-width="2"
|
|
293
|
-
d="M10 14l2-2m0 0l2-2m-2 2l-2-2m2 2l2 2m7-2a9 9 0 11-18 0 9 9 0 0118 0z"
|
|
294
|
-
/>
|
|
295
|
-
</svg>
|
|
296
|
-
<span></span>
|
|
297
|
-
</div>
|
|
298
|
-
|
|
299
|
-
<div class="flex justify-center">
|
|
300
|
-
<button onclick="visualize()" class="btn btn-secondary">
|
|
301
|
-
Visualize with Penzai!
|
|
302
|
-
</button>
|
|
303
|
-
</div>
|
|
304
|
-
|
|
305
|
-
<hr />
|
|
306
|
-
<div id="visualizationResult"></div>
|
|
307
|
-
</body>
|
|
308
|
-
</html>
|
|
@@ -1,147 +0,0 @@
|
|
|
1
|
-
Metadata-Version: 2.3
|
|
2
|
-
Name: statedict2pytree
|
|
3
|
-
Version: 0.5.0
|
|
4
|
-
Summary: Converts torch models into PyTrees for Equinox
|
|
5
|
-
Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
|
|
6
|
-
Requires-Python: ~=3.10
|
|
7
|
-
Requires-Dist: beartype
|
|
8
|
-
Requires-Dist: equinox>=0.11.4
|
|
9
|
-
Requires-Dist: flask
|
|
10
|
-
Requires-Dist: jax
|
|
11
|
-
Requires-Dist: jaxlib
|
|
12
|
-
Requires-Dist: jaxtyping
|
|
13
|
-
Requires-Dist: loguru
|
|
14
|
-
Requires-Dist: penzai
|
|
15
|
-
Requires-Dist: pydantic
|
|
16
|
-
Requires-Dist: torch
|
|
17
|
-
Requires-Dist: typing-extensions
|
|
18
|
-
Provides-Extra: dev
|
|
19
|
-
Requires-Dist: mkdocs; extra == 'dev'
|
|
20
|
-
Requires-Dist: nox; extra == 'dev'
|
|
21
|
-
Requires-Dist: pre-commit; extra == 'dev'
|
|
22
|
-
Requires-Dist: pytest; extra == 'dev'
|
|
23
|
-
Description-Content-Type: text/markdown
|
|
24
|
-
|
|
25
|
-
# statedict2pytree
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
The goal of this package is to simplify the conversion from PyTorch models into JAX PyTrees (which can be used e.g. in Equinox). The way this works is by putting both models side my side and aligning the weights in the right order. Then, all statedict2pytree is doing, is iterating over both lists and matching the weight matrices.
|
|
30
|
-
|
|
31
|
-
Usually, if you _declared the fields in the same order as in the PyTorch model_, you don't have to rearrange anything -- but the option is there if you need it.
|
|
32
|
-
|
|
33
|
-
(Theoretically, you can rearrange the model in any way you like - e.g. last layer as the first layer - as long as the shapes match!)
|
|
34
|
-
|
|
35
|
-
## Shape Matching? What's that?
|
|
36
|
-
|
|
37
|
-
Currently, there is no sophisticated shape matching in place. Two matrices are considered "matching" if the product of their shape match. For example:
|
|
38
|
-
|
|
39
|
-
1. (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
|
|
40
|
-
|
|
41
|
-
## Get Started
|
|
42
|
-
|
|
43
|
-
### Installation
|
|
44
|
-
|
|
45
|
-
Run
|
|
46
|
-
|
|
47
|
-
```bash
|
|
48
|
-
pip install statedict2pytree
|
|
49
|
-
|
|
50
|
-
```
|
|
51
|
-
|
|
52
|
-
### Basic Example
|
|
53
|
-
|
|
54
|
-
```python
|
|
55
|
-
import equinox as eqx
|
|
56
|
-
import jax
|
|
57
|
-
import torch
|
|
58
|
-
import statedict2pytree as s2p
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
def test_mlp():
|
|
62
|
-
in_size = 784
|
|
63
|
-
out_size = 10
|
|
64
|
-
width_size = 64
|
|
65
|
-
depth = 2
|
|
66
|
-
key = jax.random.PRNGKey(22)
|
|
67
|
-
|
|
68
|
-
class EqxMLP(eqx.Module):
|
|
69
|
-
mlp: eqx.nn.MLP
|
|
70
|
-
batch_norm: eqx.nn.BatchNorm
|
|
71
|
-
|
|
72
|
-
def __init__(self, in_size, out_size, width_size, depth, key):
|
|
73
|
-
self.mlp = eqx.nn.MLP(in_size, out_size, width_size, depth, key=key)
|
|
74
|
-
self.batch_norm = eqx.nn.BatchNorm(out_size, axis_name="batch")
|
|
75
|
-
|
|
76
|
-
def __call__(self, x, state):
|
|
77
|
-
return self.batch_norm(self.mlp(x), state)
|
|
78
|
-
|
|
79
|
-
jax_model = EqxMLP(in_size, out_size, width_size, depth, key)
|
|
80
|
-
|
|
81
|
-
class TorchMLP(torch.nn.Module):
|
|
82
|
-
def __init__(self, in_size, out_size, width_size, depth):
|
|
83
|
-
super(TorchMLP, self).__init__()
|
|
84
|
-
self.layers = torch.nn.ModuleList()
|
|
85
|
-
self.layers.append(torch.nn.Linear(in_size, width_size))
|
|
86
|
-
for _ in range(depth - 1):
|
|
87
|
-
self.layers.append(torch.nn.Linear(width_size, width_size))
|
|
88
|
-
self.layers.append(torch.nn.Linear(width_size, out_size))
|
|
89
|
-
self.batch_norm = torch.nn.BatchNorm1d(out_size)
|
|
90
|
-
|
|
91
|
-
def forward(self, x):
|
|
92
|
-
for layer in self.layers[:-1]:
|
|
93
|
-
x = torch.relu(layer(x))
|
|
94
|
-
x = self.batch_norm(self.layers[-1](x))
|
|
95
|
-
return x
|
|
96
|
-
|
|
97
|
-
torch_model = TorchMLP(in_size, out_size, width_size, depth)
|
|
98
|
-
state_dict = torch_model.state_dict()
|
|
99
|
-
s2p.start_conversion(jax_model, state_dict)
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
if __name__ == "__main__":
|
|
103
|
-
test_mlp()
|
|
104
|
-
|
|
105
|
-
```
|
|
106
|
-
|
|
107
|
-
There exists also a function called `s2p.convert` which does the actual conversion:
|
|
108
|
-
|
|
109
|
-
```python
|
|
110
|
-
|
|
111
|
-
class Field(BaseModel):
|
|
112
|
-
path: str
|
|
113
|
-
shape: tuple[int, ...]
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
class TorchField(Field):
|
|
117
|
-
pass
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
class JaxField(Field):
|
|
121
|
-
type: str
|
|
122
|
-
|
|
123
|
-
def convert(
|
|
124
|
-
jax_fields: list[JaxField],
|
|
125
|
-
torch_fields: list[TorchField],
|
|
126
|
-
pytree: PyTree,
|
|
127
|
-
state_dict: dict,
|
|
128
|
-
):
|
|
129
|
-
...
|
|
130
|
-
```
|
|
131
|
-
|
|
132
|
-
If your models already have the right "order", then you might as well use this function directly. Note that the lists `jax_fields` and `torch_fields` must have the same length and each matching entry must have the same shape!
|
|
133
|
-
|
|
134
|
-
For the full, automatic experience, use `autoconvert`:
|
|
135
|
-
|
|
136
|
-
```python
|
|
137
|
-
import statedict2pytree as s2p
|
|
138
|
-
|
|
139
|
-
my_model = Model(...)
|
|
140
|
-
state_dict = ...
|
|
141
|
-
|
|
142
|
-
model, state = s2p.autoconvert(my_model, state_dict)
|
|
143
|
-
|
|
144
|
-
```
|
|
145
|
-
|
|
146
|
-
This will however only work if your PyTree fields have been declared
|
|
147
|
-
in the same order as they appear in the state dict!
|
|
@@ -1,8 +0,0 @@
|
|
|
1
|
-
statedict2pytree/__init__.py,sha256=lXxSaFFvkhXweXp5oHSkg_dPjdp49OsF8xoqwX4d_4E,240
|
|
2
|
-
statedict2pytree/statedict2pytree.py,sha256=yLOWx1D-6tX1VjiEg_-JcYPTrg6KWAgw6waZQi1GNvA,7229
|
|
3
|
-
statedict2pytree/static/input.css,sha256=zBp60NAZ3bHTLQ7LWIugrCbOQdhiXdbDZjSLJfg6KOw,59
|
|
4
|
-
statedict2pytree/static/output.css,sha256=B0itthSyy_tduTWMyTK5sAry-W6WbeODnpQ-oOcQQng,33966
|
|
5
|
-
statedict2pytree/templates/index.html,sha256=Mbo8fFHV6kYRiBiiwayku-p-y3hUaLw_Yj3zn_cfmb0,10027
|
|
6
|
-
statedict2pytree-0.5.0.dist-info/METADATA,sha256=TOf10T0EZoPGAc0qSltZRcr8Ni7y4bHW7w3wzRVJH7A,4232
|
|
7
|
-
statedict2pytree-0.5.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
|
|
8
|
-
statedict2pytree-0.5.0.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|