Upload 02-model-soup.py

02-model-soup.py

@@ -0,0 +1,49 @@
+import torch
+from timm import create_model
+
+# 加载模型
+def load_model(model_path):
+    model = create_model('tf_efficientnet_b3_ns', num_classes=1605, pretrained=False)
+    model.load_state_dict(torch.load(model_path))
+    return model
+
+# 模型融合
+def model_soup(models, weights):
+    if len(models) != len(weights):
+        raise ValueError("Number of models and weights must match")
+    
+    # Normalize weights
+    weights = [w / sum(weights) for w in weights]
+    
+    # Initialize the fused model with the structure of the first model
+    fused_model = create_model('tf_efficientnet_b3_ns', num_classes=1605, pretrained=False)
+    fused_model_dict = fused_model.state_dict()
+    
+    for key in fused_model_dict.keys():
+        fused_model_dict[key] = sum(weight * models[i].state_dict()[key] for i, weight in enumerate(weights))
+    
+    fused_model.load_state_dict(fused_model_dict)
+    return fused_model
+
+# 模型权重路径
+model_paths = [
+    '/data/cjm/FungiCLEF2024/EfficientNet/output/trick_1.4.3/efficientnet_b3_epoch_28.pth',  # 权重 4
+    '/data/cjm/FungiCLEF2024/EfficientNet/output/trick_1.4.3.2/efficientnet_b3_epoch_28.pth',  # 权重 2.6
+    '/data/cjm/FungiCLEF2024/EfficientNet/output/trick_1.4.1/efficientnet_b3_epoch_23.pth',  # 权重 2.4
+    '/data/cjm/FungiCLEF2024/EfficientNet/output/trick_1.5.2/efficientnet_b3_epoch_21.pth', # 权重 1
+]
+
+# 加载模型
+models = [load_model(path) for path in model_paths]
+
+# 模型权重
+weights = [4, 2.6, 2.4, 1]
+
+# 进行模型融合
+fused_model = model_soup(models, weights)
+
+# 保存融合后的模型权重
+fused_model_path = '/data/cjm/FungiCLEF2024/EfficientNet/output/fused_model_soup.pth'
+torch.save(fused_model.state_dict(), fused_model_path)
+
+print(f"Fused model saved to {fused_model_path}")