Blog about developments of 150games.com, examples of game code, video games, list of tools and websites for develop games and apps, improve your assets and HTML, CSS, JS files and many more...
CODE - IA - Python - How to convert PyTorch `.pth` tensors to the SafeTensors format
import torch
from safetensors.torch import save_file
# Load your PyTorch model or state dict
pytorch_model = torch.load("your_model.pth")
# If it's a full model, you might need to get the state dict
if hasattr(pytorch_model, "state_dict"):
state_dict = pytorch_model.state_dict()
else:
state_dict = pytorch_model
# Save as safetensors format
save_file(state_dict, "converted_model.safetensors")
Explanation:
1. **Import necessary libraries**:
```python
import torch
from safetensors.torch import save_file
```
- `torch` is the PyTorch library
- `save_file` is a function from the safetensors library specifically designed for saving PyTorch tensors
2. **Load the original PyTorch model**:
```python
pytorch_model = torch.load("your_model.pth")
```
This loads your model or state dictionary from a .pth file into memory.
3. **Extract the state dictionary**:
```python
if hasattr(pytorch_model, "state_dict"):
state_dict = pytorch_model.state_dict()
else:
state_dict = pytorch_model
```
This checks if what you loaded is a full model object (which has a state_dict() method) or already a state dictionary. The state dictionary is a Python dictionary mapping parameter names to their tensor values.
4. **Save in SafeTensors format**:
```python
save_file(state_dict, "converted_model.safetensors")
```
This converts your PyTorch state dictionary to the SafeTensors format and saves it to a file.
SafeTensors is designed to be more secure and efficient than PyTorch's native format, providing protection against certain security vulnerabilities while potentially offering faster loading times.