Update README

README.md

@@ -24,10 +24,10 @@ lon_std = 0.0007030395253318959
 ```
 
 ### Instructions
-Our current best performing model is an ensemble of multiple models. To run it on hidden test data, first run all the model definitions.
+Our current best performing model is an ensemble of multiple models. To run it on hidden test data, first run the model definitions.
 
 #### Load and define models
-1. ConvNeXt
+
 ```python
 from transformers import AutoModelForImageClassification, PretrainedConfig, PreTrainedModel
 import torch
@@ -91,11 +91,6 @@ class CustomConvNeXtModel(PreTrainedModel):
     
         return model
 
-```
-
-2. ResNet
-``` python
-from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
 
 class CustomResNetConfig(PretrainedConfig):
     model_type = "custom-resnet"
@@ -166,14 +161,7 @@ class CustomResNetModel(nn.Module, PyTorchModelHubMixin):
         model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
 
         return model
-```
 
-3. EfficientNet
-``` python
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file  # Make sure to import this
-from transformers import AutoModelForImageClassification, PreTrainedModel, PretrainedConfig
-import torch.nn as nn
 
 class CustomEfficientNetConfig(PretrainedConfig):
     model_type = "custom-efficientnet"
@@ -236,14 +224,7 @@ class CustomEfficientNetModel(PreTrainedModel):
         model.load_state_dict(state_dict)
 
         return model
-```
 
-4. ViT
-```python
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file  
-from transformers import AutoModelForImageClassification, PreTrainedModel, PretrainedConfig
-import torch.nn as nn
 
 class CustomViTConfig(PretrainedConfig):
     model_type = "custom-vit"
@@ -303,8 +284,35 @@ class CustomViTModel(PreTrainedModel):
         model.load_state_dict(state_dict)
 
         return model
+
+
+# Define the WeightedEnsembleModel class
+class WeightedEnsembleModel(nn.Module):
+    def __init__(self, models, weights):
+        """
+        Initialize the ensemble model with individual models and their weights.
+        """
+        super(WeightedEnsembleModel, self).__init__()
+        self.models = nn.ModuleList(models)  # Wrap models in ModuleList
+        self.weights = weights
+
+    def forward(self, images):
+        """
+        Forward pass for the ensemble model.
+        Performs weighted averaging of logits from individual models.
+        """
+        ensemble_logits = torch.zeros((images.size(0), 2)).to(images.device)  # Initialize logits
+        for model, weight in zip(self.models, self.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
+
+
+
 ```
 
+
 Now, load the model weights from huggingface.
 ```python
 from transformers import AutoModelForImageClassification
@@ -312,6 +320,8 @@ import torch
 from sklearn.metrics import mean_absolute_error, mean_squared_error
 import matplotlib.pyplot as plt
 import numpy as np
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ```
 
 ```python
@@ -320,32 +330,30 @@ resnet = CustomResNetModel.from_pretrained(
     "final-project-5190/model-resnet-50-base",
     model_name="microsoft/resnet-50"
 )
-```
 
-
-```python
 #convnext
 convnext=CustomConvNeXtModel.from_pretrained(
     "final-project-5190/model-convnext-tiny-reducePlateau", 
     model_name="facebook/convnext-tiny-224")
-```
 
-```python
 #vit
 vit = CustomViTModel.from_pretrained(
     "final-project-5190/model-ViT-base",
     model_name="google/vit-base-patch16-224"
 )
-```
 
-```python
 #efficientnet
 efficientnet = CustomEfficientNetModel.from_pretrained(
     "final-project-5190/model-efficientnet-b0-base",
     model_name="google/efficientnet-b0"
 )
+
+models = [convnext, resnet, vit, efficientnet]
+weights = [0.28, 0.26, 0.20, 0.27]
 ```
 
+
+
 #### For data loading 
 ```python
 # Download
@@ -422,34 +430,6 @@ 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
@@ -472,24 +452,20 @@ rel_dataloader = DataLoader(rel_dataset, batch_size=32, shuffle=False)
 
 
 ```python
-# Release
+# ensemble
+ensemble_model = WeightedEnsembleModel(models=models, weights=weights).to(device)
 
-# Initialize lists to store predictions and actual values
+# Validation
 all_preds = []
 all_actuals = []
 
-# Move models to device and set them to evaluation mode
-for model in models:
-    model.to(device)
-    model.eval()
-
-# Perform inference on release dataset
+ensemble_model.eval()
 with torch.no_grad():
     for images, gps_coords in rel_dataloader:
         images, gps_coords = images.to(device), gps_coords.to(device)
 
-        # Weighted ensemble prediction
-        ensemble_logits = weighted_ensemble_predict(models, weights, images)
+        # Weighted ensemble prediction using the new model
+        ensemble_logits = ensemble_model(images)
 
         # Denormalize predictions and actual values
         preds = ensemble_logits.cpu() * torch.tensor([lat_std, lon_std]) + torch.tensor([lat_mean, lon_mean])
@@ -506,8 +482,8 @@ all_actuals = torch.cat(all_actuals).numpy()
 mae = mean_absolute_error(all_actuals, all_preds)
 rmse = mean_squared_error(all_actuals, all_preds, squared=False)
 
-print(f'Release Dataset Mean Absolute Error: {mae}')
-print(f'Release Dataset Root Mean Squared Error: {rmse}')
+print(f'Mean Absolute Error: {mae}')
+print(f'Root Mean Squared Error: {rmse}')
 
 # Convert predictions and actuals to meters
 latitude_mean_radians = np.radians(lat_mean)  # Convert to radians for cosine
@@ -528,6 +504,7 @@ rmse_meters = mean_squared_error(all_actuals_meters, all_preds_meters, squared=F
 
 print(f"Mean Absolute Error (meters): {mae_meters:.2f}")
 print(f"Root Mean Squared Error (meters): {rmse_meters:.2f}")
+
 ```
 
 After running inference on the release test set, our results are the following.