ezelpub's picture
Update README.md
4b7c0c1 verified

Information about the Dataset

Mean Latitude: 39.95156391970743

Latitude Std: 0.0007633062105681285

Mean Longitude: -75.19148737056214

Longitude Std: 0.0007871346840888362

Model definition

class ConvNeXtGPSPredictor(nn.Module, PyTorchModelHubMixin):
    def __init__(self, model_name="facebook/convnext-tiny-224", num_outputs=2):
        super(ConvNeXtGPSPredictor, self).__init__()

        # Load the ConvNeXt backbone from Hugging Face
        self.backbone = AutoModel.from_pretrained(model_name)

        # Get feature dimension from the backbone's output
        config = AutoConfig.from_pretrained(model_name)
        feature_dim = config.hidden_sizes[-1]  # Corrected attribute for ConvNeXt

        # Define the GPS regression head
        self.gps_head = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),  # Pool to a single spatial dimension
            nn.Flatten(),  # Flatten the tensor
            nn.LayerNorm(feature_dim),
            nn.Linear(feature_dim, num_outputs)  # Directly map to 2 GPS coordinates
        )

    def forward(self, x):
        # Extract features from the backbone
        features = self.backbone(x).last_hidden_state
        
        # Pass through the GPS head
        coords = self.gps_head(features)
        return coords

    
    def save_model(self, save_path):
        self.save_pretrained(save_path)

    def push_model(self, repo_name):
        self.push_to_hub(repo_name)

How to load the model

You can simply load the model by

model = ConvNeXtGPSPredictor.from_pretrained("cis519/convNext-GPSPredictor")