What is Safetensors and how to convert .ckpt model to .safetensors

If you often download model weight file, you will often see the .safetensors file."Safetensors" is a new file format for storing tensors, which was recently introduced by Hugging Face.

What is safetensors

Safetensors is a serialization format developed by Hugging Face that is specifically designed for efficiently storing and loading large tensors. It provides a lightweight and efficient way to serialize tensors in Python, making it easier to store and load machine learning models.

Here are some key features and benefits of Safetensors:

1. Optimized for large tensors: Safetensors uses a combination of efficient serialization and compression algorithms to reduce size of large tensors, making it faster and more efficient than other serialization formats like pickle.
2. Cross-platform compatibility: Safetensors can be used to serialize tensors in Python, and the resulting files can be easily loaded in other languages and platforms, including C++, Java, and JavaScript.
3. Easy to use: Safetensors provides a simple and intuitive API that makes it easy to serialize and deserialize tensors in Python.
4. Fast serialization and deserialization: Safetensors is designed to be fast, and it can serialize and deserialize large tensors quickly, making it ideal for use in deep learning applications.
5. Safe and secure: Safetensors uses a checksum mechanism to ensure that serialized tensors are not corrupted during storage or transfer, providing an extra layer of security.

How to use Safetensors

Install the Safetensors package: You can install the Safetensors package using pip

pip install safetensors

Serialize your tensors by safetensors.dump function. Here's an example:

import torch
 import safetensors
 ​
 # create a tensor
 x = torch.tensor([1, 2, 3])
 ​
 # serialize the tensor to a Safetensors file
 safetensors.dump(x, 'mytensor.st')

This will create a file named “mytensor.st” in the current directory that contains serialized tensor.

Deserialize your tensors by safetensors.load function

import safetensors
 ​
 # deserialize the tensor from the Safetensors file
 x = safetensors.load('mytensor.st')
 ​
 # print the deserialized tensor
 print(x)

What is different with pickle

While both Safetensors and Pickle can be used for serializing and deserializing Python objects, there are some key differences:

1. Safety: Pickle is not considered to be a safe format for storing and sharing data because it can execute arbitrary code during deserialization, which could potentially be a security risk. Safetensors was designed to provide a secure format for storing tensors and models, with features such as encryption and access control.
2. Portability: Pickle is designed specifically for Python and is not always compatible with other programming languages. Safetensors is designed to be compatible with various deep learning frameworks and libraries, allowing users to share their models and data across different tools and workflows.
3. Performance: Pickle can be slow when serializing and deserializing large Python objects, especially when compared to more optimized serialization formats like Protocol Buffers or Apache Arrow. Safetensors is designed to be fast and efficient for storing and sharing tensors, which are fundamental building blocks of many deep learning models.

How to convert .ckpt model to .safetensors

Tensorflow

Load TensorFlow model: Use TensorFlow’s tf.train.Checkpoint class to load model from the checkpoint file. Here's an example:

import tensorflow as tf
 ​
 # define your model architecture
 inputs = tf.keras.Input(shape=(32,))
 x = tf.keras.layers.Dense(64, activation='relu')(inputs)
 outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
 model = tf.keras.Model(inputs=inputs, outputs=outputs)
 ​
 # load the model weights from the checkpoint file
 checkpoint_path = 'my_model.ckpt'
 ckpt = tf.train.Checkpoint(model=model)
 ckpt.restore(checkpoint_path).expect_partial()

This will load model architecture and weights from the checkpoint file.

Use Safetensors’ safetensors.dump function to serialize model to Safetensors file:

import safetensors
 ​
 # convert the model to Safetensors format and save it to a file
 safetensors.dump(model, 'my_model.st')

This will convert your TensorFlow model to the Safetensors format and save it to a file named “my_model.st” in current directory.

Pytorch

Use PyTorch’s torch.load function to load model from saved checkpoint file:

import torch
 ​
 # define your model architecture
 class MyModel(torch.nn.Module):
     def __init__(self):
         super(MyModel, self).__init__()
         self.fc1 = torch.nn.Linear(32, 64)
         self.fc2 = torch.nn.Linear(64, 10)
 ​
     def forward(self, x):
         x = self.fc1(x)
         x = torch.nn.functional.relu(x)
         x = self.fc2(x)
         x = torch.nn.functional.softmax(x, dim=1)
         return x
 ​
 # load the model weights from the checkpoint file
 checkpoint_path = 'my_model.pt'
 model = MyModel()
 model.load_state_dict(torch.load(checkpoint_path))

This will load your model architecture and the weights from the saved checkpoint file. Use Safetensors’ safetensors.dump function to serialize model to Safetensors file:

import safetensors
 ​
 # convert the model to Safetensors format and save it to a file
 safetensors.dump(model, 'my_model.st')

This will convert your PyTorch model to Safetensors format and save it to a file named “my_model.st” in the current directory.

Bonus:Convert tensorflow weight to pytorch with Safetensors

Safetensors can convert TensorFlow weights to PyTorch or convert PyTorch weights to TensorFlow

 import tensorflow as tf
 import safetensors
 ​
 # define your TensorFlow model architecture
 inputs = tf.keras.Input(shape=(32,))
 x = tf.keras.layers.Dense(64, activation='relu')(inputs)
 outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
 model = tf.keras.Model(inputs=inputs, outputs=outputs)
 ​
 # load the model weights from the checkpoint file
 checkpoint_path = 'my_model.ckpt'
 ckpt = tf.train.Checkpoint(model=model)
 ckpt.restore(checkpoint_path).expect_partial()
 ​
 # convert the TensorFlow weights to Safetensors format
 safetensors.dump(ckpt.model, 'my_weights.st')

This will convert your TensorFlow weights to the Safetensors format and save them to a file named “my_weights.st”.

And next we can load Safetensors weights in PyTorch:

 import torch
 ​
 # define your PyTorch model architecture
 class MyModel(torch.nn.Module):
     def __init__(self):
         super(MyModel, self).__init__()
         self.fc1 = torch.nn.Linear(32, 64)
         self.fc2 = torch.nn.Linear(64, 10)
 ​
     def forward(self, x):
         x = self.fc1(x)
         x = torch.nn.functional.relu(x)
         x = self.fc2(x)
         x = torch.nn.functional.softmax(x, dim=1)
         return x
 ​
 # load the Safetensors weights into the PyTorch model
 weights_path = 'my_weights.st'
 model = MyModel()
 model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))

This will load the Safetensors weights into your PyTorch model. You can then use this model for inference.

Note : this method only converts the weights, not entire model architecture, so you’ll need to define a corresponding PyTorch model architecture that matches TensorFlow model architecture.