训练权重匹配策略

rom torch import nn
import torch
import os#环境代理设置
import torchvision.models as models
os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"
model=timm.create_model('vit_small_patch16_224.dino',pretrained=True,in_chans=3,img_size=(1280,1280))
#m=nn.Sequential(*(m.children()))
class CustomConvit(nn.Module):
    def __init__(self):
        super().__init__()
        # 直接将Convit模型的基础部分添加到Sequential中
        self.features = model
    def forward(self,x):
        x = self.features(x)
        return x
data=torch.ones(1,3,1280,1280)
# 创建自定义的Convit模型
m= CustomConvit()
print(m(data).shape)
# 打印模型的参数字典的键
print("Model State Dict Keys:")
for key in m.features.state_dict().keys():
    print(key,end='')
    print((m.features.state_dict().get(key)).shape)
# 加载你自定义的权重
checkpoint_path = 'vit_small_patch16_224.dino.bin'
checkpoint = torch.load(checkpoint_path)
#rint(checkpoint)
# 打印自定义权重的键
print("\nCustom Weights Dict Keys:")
for key in checkpoint.keys():
    print(key,end='')
    print((m.features.state_dict().get(key)).shape)
# 创建新的权重字典，仅包含相等的键
matching_weights = {key: checkpoint[key] for key in m.features.state_dict().keys() if key in checkpoint and checkpoint[key].shape==m.features.state_dict().get(key).shape}

# 打印匹配的键的信息
for key, value in matching_weights.items():
    print(f"Key: {key}, Shape: {value.shape}")

# 加载匹配的权重到模型
model.load_state_dict(matching_weights, strict=False)