safetensor的使用_safetensor加载checkpoint

Huggingface更新之后保存的张量变成了safetensor，它和bin文件很相似

import safetensors.torch
checkpoint=safetensors.torch.load_file("path/to/model.safetensor")
1
2

这里的checkpoint就是和bin加载时一样的字典数据类型，模型如何加载和bin文件一样。

由于需要经常配置baseline，所以需要经常替换权重

在这里，我一般直接使用字典加载权重，多余的部分直接删除掉，公有的，最核心的内容，直接字典加载

checkpoint=torch.load(self.pretrained_path)
# or
checkpoint=safetensors.torch.load_file("path/to/model.safetensor")

for k, v in checkpoint.items():
    if k.startswith('img_encoder.model'):
        new_key = k[len('img_encoder.model.'):]
        tmp_dict.update({new_key:checkpoint[k]})                
         
 model_dict=self.visual_extractor.state_dict()
 unique_keys = set(model_dict) ^ set(tmp_dict)
 for k in unique_keys:
     del tmp_dict[k] 

 self.visual_extractor.load_state_dict(tmp_dict)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

这里，startwith目的就是去掉前缀等信息，提取正确的组件。