O que é SafeTensor?
O SafeTensor é um formato de serialização de tensores criado para oferecer uma alternativa segura e eficiente ao tradicional formato PyTorch `.pt` ou `.pth`. Desenvolvido inicialmente por pesquisadores da Hugging Face, esse novo padrão de salvamento de pesos de modelos de machine learning visa eliminar riscos de segurança associados ao uso de pickle — mecanismo de serialização do Python que, apesar de poderoso, pode ser explorado para execução de código malicioso.
Se você trabalha com inteligência artificial, aprendizado de máquina ou deep learning, especialmente em ambientes com múltiplos colaboradores, como notebooks colaborativos ou servidores compartilhados, compreender o que é SafeTensor e como utilizá-lo pode melhorar tanto a segurança quanto a portabilidade dos seus projetos.
Por que o formato SafeTensor foi criado?
O principal motivo para a criação do SafeTensor foi a insegurança presente nos formatos tradicionais de salvamento de modelos, como os arquivos `.pt` do PyTorch, que utilizam o módulo pickle para serializar objetos Python. O pickle, embora eficiente e flexível, permite a execução arbitrária de código Python no momento do carregamento de arquivos. Isso significa que um arquivo `.pt` pode conter código malicioso que será executado automaticamente ao ser carregado, colocando em risco sistemas e dados.
O SafeTensor resolve esse problema adotando uma estrutura simples baseada em JSON para metadados e uma matriz binária compactada que representa os tensores. Esse formato garante que os arquivos possam ser carregados de forma segura, sem a execução de nenhum código externo, sendo, portanto, imunes a ataques via execução de código malicioso embutido.
Vantagens do SafeTensor
- Segurança: Não executa código ao ser carregado, evitando riscos de execução de scripts maliciosos.
- Eficiência: Armazenamento binário direto e leitura rápida dos dados.
- Compatibilidade: Suporte crescente em bibliotecas como PyTorch, TensorFlow (via conversão) e Transformers da Hugging Face.
- Portabilidade: Pode ser transferido entre sistemas com diferentes versões de bibliotecas ou plataformas, sem riscos.
- Leitura parcial: Permite leitura seletiva de variáveis/tensores dentro do arquivo, útil para modelos grandes.
Estrutura interna do SafeTensor
O formato SafeTensor é composto por duas partes principais:
- Header (cabeçalho): Contém metadados em formato JSON, descrevendo cada tensor, seu dtype, forma (shape), offset e comprimento.
- Dados binários: Todos os tensores são armazenados sequencialmente como um buffer binário contínuo.
Essa estrutura torna o arquivo simples e direto, evitando estruturas complexas como grafos de objetos (usados pelo pickle), o que também facilita o parse e análise de arquivos mesmo sem bibliotecas especializadas.
Comparação entre SafeTensor e .pt/.pth
| Característica | .pt/.pth (pickle) | SafeTensor |
|---|---|---|
| Segurança | Inseguro (executa código) | Seguro (sem execução) |
| Formato | Binário + pickle | JSON + binário |
| Leitura parcial | Não | Sim |
| Compatibilidade | PyTorch | Hugging Face, PyTorch (limitado) |
| Velocidade | Boa | Melhor |
Como usar o SafeTensor na prática
Para começar a usar o formato SafeTensor, é necessário instalar a biblioteca oficial:
pip install safetensorsSalvando tensores com SafeTensor
import torch
from safetensors.torch import save_file
# Criar alguns tensores
data = {
"peso1": torch.randn(3, 3),
"peso2": torch.ones(2, 2)
}
# Salvar no formato SafeTensor
save_file(data, "modelo.safetensors")Carregando tensores com SafeTensor
from safetensors.torch import load_file
data_carregado = load_file("modelo.safetensors")
print(data_carregado["peso1"])SafeTensor na Hugging Face
A Hugging Face tem adotado o SafeTensor como padrão para o armazenamento de modelos no seu hub. Isso significa que, ao baixar modelos recentes como o BERT, GPT-NeoX, ou LLaMA2, é possível encontrar versões já no formato `.safetensors`.
Além disso, a biblioteca `transformers` já inclui suporte nativo ao formato, bastando garantir que o arquivo tenha a extensão correta e que o pacote `safetensors` esteja instalado.
from transformers import AutoModel
model = AutoModel.from_pretrained("nome-do-modelo", trust_remote_code=True)Se o modelo tiver arquivos `.safetensors`, eles serão utilizados automaticamente se o pacote estiver presente.
Limitações e cuidados
Apesar das inúmeras vantagens, o formato SafeTensor também possui algumas limitações:
- Não serializa objetos Python complexos (apenas tensores).
- Incompatível com modelos salvos por `torch.save(model)` diretamente.
- Necessita da biblioteca `safetensors` para leitura e escrita.
Portanto, é recomendado seu uso em ambientes onde os modelos já são serializados com os pesos em dicionários (`state_dict()`), o que é uma prática comum em ambientes profissionais.
Como converter modelos PyTorch existentes para SafeTensor
Se você já possui modelos treinados em formato `.pt` ou `.pth`, é possível convertê-los para SafeTensor com poucas linhas de código:
import torch
from safetensors.torch import save_file
state_dict = torch.load("modelo_antigo.pt", map_location="cpu")
save_file(state_dict, "modelo_convertido.safetensors")Exemplos de uso no mundo real
Diversos projetos open-source e iniciativas da comunidade já adotam o SafeTensor:
- Hugging Face Transformers: Muitos modelos vêm com `.safetensors` por padrão.
- Spaces com Gradio: Usam SafeTensor para carregar modelos com segurança.
- Serviços de deploy de IA: SafeTensor é preferido para ambientes de produção.
FAQ sobre SafeTensor
SafeTensor é mais rápido que pickle? Sim, o carregamento é mais rápido por causa do acesso direto à memória e à estrutura binária contínua.
Posso usar SafeTensor com modelos TensorFlow? Não diretamente, mas pode-se converter os pesos para NumPy e depois salvar com a versão `safetensors.numpy`.
Qual o tamanho dos arquivos SafeTensor? Geralmente comparável ou menor, dependendo da compressão implícita e estrutura dos tensores.
É obrigatório usar SafeTensor? Não, mas é altamente recomendado, principalmente para quem publica modelos ou compartilha com terceiros.
Conclusão
O SafeTensor é uma revolução silenciosa na forma como armazenamos pesos de modelos de machine learning. Ao eliminar os riscos associados à execução de código e ao mesmo tempo manter performance, compatibilidade e simplicidade, ele se posiciona como o novo padrão seguro para modelos em produção.
Se você ainda utiliza o formato `.pt` para compartilhar seus modelos, considere migrar para o SafeTensor. A mudança é simples e pode evitar falhas de segurança catastróficas.
Com o suporte crescente da comunidade e ferramentas cada vez mais otimizadas, o SafeTensor está preparado para ser o novo padrão da IA moderna.