Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import torch.nn as nn
|
|
|
|
|
|
class ConvAP(nn.Module):
|
|
"""Implementation of ConvAP as of https://arxiv.org/pdf/2210.10239.pdf
|
|
|
|
Args:
|
|
in_channels (int): number of channels in the input of ConvAP
|
|
out_channels (int, optional): number of channels that ConvAP outputs. Defaults to 512.
|
|
s1 (int, optional): spatial height of the adaptive average pooling. Defaults to 2.
|
|
s2 (int, optional): spatial width of the adaptive average pooling. Defaults to 2.
|
|
"""
|
|
def __init__(self, in_channels, out_channels=512, s1=2, s2=2):
|
|
super(ConvAP, self).__init__()
|
|
self.channel_pool = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=True)
|
|
self.AAP = nn.AdaptiveAvgPool2d((s1, s2))
|
|
|
|
def forward(self, x):
|
|
#
|
|
x, t = x #dinov2专属
|
|
# x = self.channel_pool(x)
|
|
x = self.AAP(x)
|
|
x = F.normalize(x.flatten(1), p=2, dim=1)
|
|
return x
|
|
|
|
|
|
if __name__ == '__main__':
|
|
x = torch.randn(4, 2048, 10, 10)
|
|
m = ConvAP(2048, 512)
|
|
r = m(x)
|
|
print(r.shape) |