Modified readme

README.md

@@ -422,6 +422,34 @@ inference_transform = transforms.Compose([
 ])
 ```
 
+### Ensemble
+Define Ensemble (weighted average) and prepare model
+```python
+models = [convnext, resnet, vit, efficientnet]
+weights = [0.28, 0.26, 0.20, 0.27] # based on val 1/RMSE
+```
+
+```python
+# Weighted ensemble prediction function
+def weighted_ensemble_predict(models, weights, images):
+    """
+    Generate weighted ensemble predictions by averaging logits using model weights.
+    """
+    ensemble_logits = torch.zeros((images.size(0), 2)).to(images.device)  # Initialize logits for ensemble
+    for model, weight in zip(models, weights):
+        outputs = model(images)
+        logits = outputs.logits if hasattr(outputs, "logits") else outputs  # Extract logits
+        ensemble_logits += weight * logits  # Weighted sum of logits
+    return ensemble_logits  # Return the weighted logits sum (no division since weights sum to 1)
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+for model in models:
+    model.to(device)
+    model.eval()
+
+```
+
 Here's an exmaple of us testing the ensemble on the release test set. You can just change the load release_data line below and run the rest of the code to obtain rMSE.
 
 ```python
@@ -442,10 +470,6 @@ rel_dataset = GPSImageDataset(
 rel_dataloader = DataLoader(rel_dataset, batch_size=32, shuffle=False)
 ```
 
-```python
-models = [convnext, resnet, vit, efficientnet]
-weights = [0.28, 0.26, 0.20, 0.27] # based on val 1/RMSE
-```
 
 ```python
 # Release