File size: 1,622 Bytes
cceb9a3 4b7c0c1 cceb9a3 4b7c0c1 cceb9a3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
# Information about the Dataset
**Mean Latitude**: 39.95156391970743
**Latitude Std**: 0.0007633062105681285
**Mean Longitude**: -75.19148737056214
**Longitude Std**: 0.0007871346840888362
# Model definition
```python
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
```python
model = ConvNeXtGPSPredictor.from_pretrained("cis519/convNext-GPSPredictor")
``` |