|
|
|
|
|
import pickle |
|
import torch |
|
from torch import nn |
|
|
|
from detectron2.utils.file_io import PathManager |
|
|
|
from .utils import normalize_embeddings |
|
|
|
|
|
class VertexDirectEmbedder(nn.Module): |
|
""" |
|
Class responsible for embedding vertices. Vertex embeddings take |
|
the form of a tensor of size [N, D], where |
|
N = number of vertices |
|
D = number of dimensions in the embedding space |
|
""" |
|
|
|
def __init__(self, num_vertices: int, embed_dim: int): |
|
""" |
|
Initialize embedder, set random embeddings |
|
|
|
Args: |
|
num_vertices (int): number of vertices to embed |
|
embed_dim (int): number of dimensions in the embedding space |
|
""" |
|
super(VertexDirectEmbedder, self).__init__() |
|
self.embeddings = nn.Parameter(torch.Tensor(num_vertices, embed_dim)) |
|
self.reset_parameters() |
|
|
|
@torch.no_grad() |
|
def reset_parameters(self): |
|
""" |
|
Reset embeddings to random values |
|
""" |
|
self.embeddings.zero_() |
|
|
|
def forward(self) -> torch.Tensor: |
|
""" |
|
Produce vertex embeddings, a tensor of shape [N, D] where: |
|
N = number of vertices |
|
D = number of dimensions in the embedding space |
|
|
|
Return: |
|
Full vertex embeddings, a tensor of shape [N, D] |
|
""" |
|
return normalize_embeddings(self.embeddings) |
|
|
|
@torch.no_grad() |
|
def load(self, fpath: str): |
|
""" |
|
Load data from a file |
|
|
|
Args: |
|
fpath (str): file path to load data from |
|
""" |
|
with PathManager.open(fpath, "rb") as hFile: |
|
data = pickle.load(hFile) |
|
for name in ["embeddings"]: |
|
if name in data: |
|
getattr(self, name).copy_( |
|
torch.tensor(data[name]).float().to(device=getattr(self, name).device) |
|
) |
|
|
