File size: 515 Bytes
0f21852 |
1 2 3 4 5 6 7 8 9 10 11 12 13 |
# This script exports pre-trained model weights in the safetensors format.
import numpy as np
import torch
import torchvision
from safetensors import torch as stt
m = torchvision.models.efficientnet_b5(pretrained=True)
stt.save_file(m.state_dict(), 'efficientnet-b5.safetensors')
m = torchvision.models.efficientnet_b6(pretrained=True)
stt.save_file(m.state_dict(), 'efficientnet-b6.safetensors')
m = torchvision.models.efficientnet_b7(pretrained=True)
stt.save_file(m.state_dict(), 'efficientnet-b7.safetensors')
|