Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
17 lines
603 B
Python
17 lines
603 B
Python
import torch
|
|
import torch.nn.functional as F
|
|
import torch.nn as nn
|
|
|
|
class GeMPool(nn.Module):
|
|
"""Implementation of GeM as in https://github.com/filipradenovic/cnnimageretrieval-pytorch
|
|
we add flatten and norm so that we can use it as one aggregation layer.
|
|
"""
|
|
def __init__(self, p=3, eps=1e-6):
|
|
super().__init__()
|
|
self.p = nn.Parameter(torch.ones(1)*p)
|
|
self.eps = eps
|
|
|
|
def forward(self, x):
|
|
x = F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1./self.p)
|
|
x = x.flatten(1)
|
|
return F.normalize(x, p=2, dim=1) |