Views
No views yet
1class ConvNeXtGPSPredictor(nn.Module, PyTorchModelHubMixin):
2 def __init__(self, model_name="facebook/convnext-tiny-224", num_outputs=2):
3 super(ConvNeXtGPSPredictor, self).__init__()
4
5 # Load the ConvNeXt backbone from Hugging Face
6 self.backbone = AutoModel.from_pretrained(model_name)
7
8 # Get feature dimension from the backbone's output
9 config = AutoConfig.from_pretrained(model_name)
10 feature_dim = config.hidden_sizes[-1] # Corrected attribute for ConvNeXt
11
12 # Define the GPS regression head
13 self.gps_head = nn.Sequential(
14 nn.AdaptiveAvgPool2d((1, 1)), # Pool to a single spatial dimension
15 nn.Flatten(), # Flatten the tensor
16 nn.LayerNorm(feature_dim),
17 nn.Linear(feature_dim, num_outputs) # Directly map to 2 GPS coordinates
18 )
19
20 def forward(self, x):
21 # Extract features from the backbone
22 features = self.backbone(x).last_hidden_state
23
24 # Pass through the GPS head
25 coords = self.gps_head(features)
26 return coords
27
28
29 def save_model(self, save_path):
30 self.save_pretrained(save_path)
31
32 def push_model(self, repo_name):
33 self.push_to_hub(repo_name)model = ConvNeXtGPSPredictor.from_pretrained("cis519/convNext-GPSPredictor")