Conversation
| PRECISIONS = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16} | ||
|
|
||
|
|
||
| def flatten_for_inspection(tree, parent_key="", separator="/"): |
There was a problem hiding this comment.
nit: you could use flax.nnx.traversals.flatten_mapping here
| import jax.numpy as jnp | ||
| import numpy as np | ||
| import orbax.checkpoint as ocp | ||
| from safetensors.torch import load_model |
There was a problem hiding this comment.
nit: import by module, rather than by function
| return nnx.merge(graphdef, state) | ||
|
|
||
| def load_pytorch(self, train_config, weight_path: str): | ||
| print(f"train_config: {train_config}") |
There was a problem hiding this comment.
generally shouldn't print from production code, can use logger.info if you want
| output_transforms: Sequence[_transforms.DataTransformFn] = (), | ||
| sample_kwargs: dict[str, Any] | None = None, | ||
| metadata: dict[str, Any] | None = None, | ||
| device: str = "cpu", |
There was a problem hiding this comment.
maybe call this torch_device to make it clear that it only affects PyTorch?
| sample_kwargs=sample_kwargs, | ||
| metadata=train_config.policy_metadata, | ||
| is_pytorch=is_pytorch, | ||
| device="cuda" if is_pytorch else None, |
There was a problem hiding this comment.
actually, seems like maybe the device isn't needed since you hardcode to "cuda"
| T = TypeVar("T") | ||
|
|
||
|
|
||
| # Custom type checker that handles both JAX arrays and torch tensors |
There was a problem hiding this comment.
I'm confused why this is necessary -- it doesn't do anything special
|
|
||
| jaxtyping._decorator._check_dataclass_annotations = _check_dataclass_annotations # noqa: SLF001 | ||
|
|
||
| TorchTensor = torch.Tensor |
There was a problem hiding this comment.
I think that, instead of all this, you can just redefine Array = jax.Array | torch.Tensor and it should work anywhere. And then you don't need to change the usage either, you can just do at.Float[Array, "b ..."]. jaxtyping is supposed to have built-in PyTorch support.
698ed43 to
f83909e
Compare
| 1. Upgrade the transformers library to 4.53.2 | ||
| - The required version is already specified in pyproject.toml | ||
| - If you set up your environment previously, reinstall it to ensure you have transformers 4.53.2 | ||
| - You can verify the version with `pip show transformers` |
There was a problem hiding this comment.
| - You can verify the version with `pip show transformers` | |
| - You can verify the version with `uv pip show transformers` |
2 participants
This PR implements pytorch model under openpi. Currently it supports inference only for pi0 and pi05 models. Finetuning WIP.
I have done the following tests.
TODOs:
pytest .to make sure no tests are broken