File size: 5,217 Bytes
cb851a6 |
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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import torch
import torch.nn as nn
from typing import Optional
from torch import Tensor
import numpy as np
from huggingface_hub import PyTorchModelHubMixin
# Constants
A1 = 1.340264
A2 = -0.081106
A3 = 0.000893
A4 = 0.003796
SF = 66.50336
@torch.jit.script
def gaussian_encoding(
v: Tensor,
b: Tensor) -> Tensor:
r"""Computes :math:`\gamma(\mathbf{v}) = (\cos{2 \pi \mathbf{B} \mathbf{v}} , \sin{2 \pi \mathbf{B} \mathbf{v}})`
Args:
v (Tensor): input tensor of shape :math:`(N, *, \text{input_size})`
b (Tensor): projection matrix of shape :math:`(\text{encoded_layer_size}, \text{input_size})`
Returns:
Tensor: mapped tensor of shape :math:`(N, *, 2 \cdot \text{encoded_layer_size})`
See :class:`~rff.layers.GaussianEncoding` for more details.
"""
vp = 2 * np.pi * v @ b.T
return torch.cat((torch.cos(vp), torch.sin(vp)), dim=-1)
def sample_b(sigma: float, size: tuple) -> Tensor:
r"""Matrix of size :attr:`size` sampled from from :math:`\mathcal{N}(0, \sigma^2)`
Args:
sigma (float): standard deviation
size (tuple): size of the matrix sampled
See :class:`~rff.layers.GaussianEncoding` for more details
"""
return torch.randn(size) * sigma
class GaussianEncoding(nn.Module):
"""Layer for mapping coordinates using random Fourier features"""
def __init__(self, sigma: Optional[float] = None,
input_size: Optional[float] = None,
encoded_size: Optional[float] = None,
b: Optional[Tensor] = None):
r"""
Args:
sigma (Optional[float]): standard deviation
input_size (Optional[float]): the number of input dimensions
encoded_size (Optional[float]): the number of dimensions the `b` matrix maps to
b (Optional[Tensor], optional): Optionally specify a :attr:`b` matrix already sampled
Raises:
ValueError:
If :attr:`b` is provided and one of :attr:`sigma`, :attr:`input_size`,
or :attr:`encoded_size` is provided. If :attr:`b` is not provided and one of
:attr:`sigma`, :attr:`input_size`, or :attr:`encoded_size` is not provided.
"""
super().__init__()
if b is None:
if sigma is None or input_size is None or encoded_size is None:
raise ValueError(
'Arguments "sigma," "input_size," and "encoded_size" are required.')
b = sample_b(sigma, (encoded_size, input_size))
elif sigma is not None or input_size is not None or encoded_size is not None:
raise ValueError('Only specify the "b" argument when using it.')
self.b = nn.parameter.Parameter(b, requires_grad=False)
def forward(self, v: Tensor) -> Tensor:
r"""Computes :math:`\gamma(\mathbf{v}) = (\cos{2 \pi \mathbf{B} \mathbf{v}} , \sin{2 \pi \mathbf{B} \mathbf{v}})`
Args:
v (Tensor): input tensor of shape :math:`(N, *, \text{input_size})`
Returns:
Tensor: Tensor mapping using random fourier features of shape :math:`(N, *, 2 \cdot \text{encoded_size})`
"""
return gaussian_encoding(v, self.b)
def equal_earth_projection(L):
latitude = L[:, 0]
longitude = L[:, 1]
latitude_rad = torch.deg2rad(latitude)
longitude_rad = torch.deg2rad(longitude)
sin_theta = (torch.sqrt(torch.tensor(3.0)) / 2) * torch.sin(latitude_rad)
theta = torch.asin(sin_theta)
denominator = 3 * (9 * A4 * theta**8 + 7 * A3 * theta**6 + 3 * A2 * theta**2 + A1)
x = (2 * torch.sqrt(torch.tensor(3.0)) * longitude_rad * torch.cos(theta)) / denominator
y = A4 * theta**9 + A3 * theta**7 + A2 * theta**3 + A1 * theta
return (torch.stack((x, y), dim=1) * SF) / 180
class LocationEncoderCapsule(nn.Module):
def __init__(self, sigma):
super(LocationEncoderCapsule, self).__init__()
rff_encoding = GaussianEncoding(sigma=sigma, input_size=2, encoded_size=256)
self.km = sigma
self.capsule = nn.Sequential(rff_encoding,
nn.Linear(512, 1024),
nn.ReLU(),
nn.Linear(1024, 1024),
nn.ReLU(),
nn.Linear(1024, 1024),
nn.ReLU())
self.head = nn.Sequential(nn.Linear(1024, 512))
def forward(self, x):
x = self.capsule(x)
x = self.head(x)
return x
class LocationEncoder(nn.Module, PyTorchModelHubMixin):
def __init__(self):
super(LocationEncoder, self).__init__()
self.sigma = [2**0, 2**4, 2**8]
self.n = len(self.sigma)
for i, s in enumerate(self.sigma):
self.add_module('LocEnc' + str(i), LocationEncoderCapsule(sigma=s))
def forward(self, location):
location = equal_earth_projection(location)
location_features = torch.zeros(location.shape[0], 512).to(location.device)
for i in range(self.n):
location_features += self._modules['LocEnc' + str(i)](location)
return location_features |