|
- #
- # Copyright (C) Stanislaw Adaszewski, 2020
- # License: GPLv3
- #
-
-
- import numpy as np
- import torch
- import torch.utils.data
- from typing import List, \
- Union, \
- Tuple
- from .data import Data, \
- EdgeType
- from .cumcount import cumcount
- import time
-
-
- def fixed_unigram_candidate_sampler(
- true_classes: torch.Tensor,
- num_repeats: torch.Tensor,
- unigrams: torch.Tensor,
- distortion: float = 1.) -> torch.Tensor:
-
- if len(true_classes.shape) != 2:
- raise ValueError('true_classes must be a 2D matrix with shape (num_samples, num_true)')
-
- if len(num_repeats.shape) != 1:
- raise ValueError('num_repeats must be 1D')
-
- num_rows = true_classes.shape[0]
- print('true_classes.shape:', true_classes.shape)
- # unigrams = np.array(unigrams)
- if distortion != 1.:
- unigrams = unigrams.to(torch.float64) ** distortion
- print('unigrams:', unigrams)
-
- indices = torch.arange(num_rows)
- indices = torch.repeat_interleave(indices, num_repeats)
- indices = torch.cat([ torch.arange(len(indices)).view(-1, 1),
- indices.view(-1, 1) ], dim=1)
-
- num_samples = len(indices)
- result = torch.zeros(num_samples, dtype=torch.long)
- print('num_rows:', num_rows, 'num_samples:', num_samples)
-
- while len(indices) > 0:
- print('len(indices):', len(indices))
- print('indices:', indices)
- sampler = torch.utils.data.WeightedRandomSampler(unigrams, len(indices))
- candidates = torch.tensor(list(sampler))
- candidates = candidates.view(len(indices), 1)
- print('candidates:', candidates)
- print('true_classes:', true_classes[indices[:, 1], :])
- result[indices[:, 0]] = candidates.transpose(0, 1)
- print('result:', result)
- mask = (candidates == true_classes[indices[:, 1], :])
- mask = mask.sum(1).to(torch.bool)
- print('mask:', mask)
- indices = indices[mask]
- # result[indices] = 0
- return result
-
-
- def get_edges_and_degrees(adj_mat: torch.Tensor) -> \
- Tuple[torch.Tensor, torch.Tensor]:
-
- if adj_mat.is_sparse:
- adj_mat = adj_mat.coalesce()
- degrees = torch.zeros(adj_mat.shape[1], dtype=torch.int64,
- device=adj_mat.device)
- degrees = degrees.index_add(0, adj_mat.indices()[1],
- torch.ones(adj_mat.indices().shape[1], dtype=torch.int64,
- device=adj_mat.device))
- edges_pos = adj_mat.indices().transpose(0, 1)
- else:
- degrees = adj_mat.sum(0)
- edges_pos = torch.nonzero(adj_mat, as_tuple=False)
- return edges_pos, degrees
-
-
- def get_true_classes(adj_mat: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- indices = adj_mat.indices()
- row_count = torch.zeros(adj_mat.shape[0], dtype=torch.long)
- #print('indices[0]:', indices[0], count[indices[0]])
- row_count = row_count.index_add(0, indices[0],
- torch.ones(indices.shape[1], dtype=torch.long))
- #print('count:', count)
- max_true_classes = torch.max(row_count).item()
- #print('max_true_classes:', max_true_classes)
- true_classes = torch.full((adj_mat.shape[0], max_true_classes),
- -1, dtype=torch.long)
-
-
- # inv = torch.unique(indices[0], return_inverse=True)
-
- # indices = indices.copy()
- # true_classes[indices[0], 0] = indices[1]
- t = time.time()
- cc = cumcount(indices[0].cpu().numpy())
- print('cumcount() took:', time.time() - t)
- cc = torch.tensor(cc)
- t = time.time()
- true_classes[indices[0], cc] = indices[1]
- print('assignment took:', time.time() - t)
-
- ''' count = torch.zeros(adj_mat.shape[0], dtype=torch.long)
- for i in range(indices.shape[1]):
- # print('looping...')
- row = indices[0, i]
- col = indices[1, i]
- #print('row:', row, 'col:', col, 'count[row]:', count[row])
- true_classes[row, count[row]] = col
- count[row] += 1 '''
-
- # t = time.time()
- # true_classes = torch.repeat_interleave(true_classes, row_count, dim=0)
- # print('repeat_interleave() took:', time.time() - t)
-
- return true_classes, row_count
-
-
- def negative_sample_adj_mat(adj_mat: torch.Tensor) -> torch.Tensor:
- if not isinstance(adj_mat, torch.Tensor):
- raise ValueError('adj_mat must be a torch.Tensor, got: %s' % adj_mat.__class__.__name__)
-
- edges_pos, degrees = get_edges_and_degrees(adj_mat)
-
- true_classes, row_count = get_true_classes(adj_mat)
- # true_classes = edges_pos[:, 1].view(-1, 1)
- # print('true_classes:', true_classes)
-
- neg_neighbors = fixed_unigram_candidate_sampler(
- true_classes, row_count, degrees, 0.75).to(adj_mat.device)
-
- print('neg_neighbors:', neg_neighbors)
-
- edges_neg = torch.cat([ edges_pos[:, 0].view(-1, 1),
- neg_neighbors.view(-1, 1) ], 1)
-
- adj_mat_neg = torch.sparse_coo_tensor(indices = edges_neg.transpose(0, 1),
- values=torch.ones(len(edges_neg)), size=adj_mat.shape,
- dtype=adj_mat.dtype, device=adj_mat.device)
-
- adj_mat_neg = adj_mat_neg.coalesce()
- indices = adj_mat_neg.indices()
- adj_mat_neg = torch.sparse_coo_tensor(indices,
- torch.ones(indices.shape[1]), adj_mat.shape,
- dtype=adj_mat.dtype, device=adj_mat.device)
-
- adj_mat_neg = adj_mat_neg.coalesce()
-
- return adj_mat_neg
-
-
- def negative_sample_data(data: Data) -> Data:
- new_edge_types = {}
- res = Data(target_value=0)
- for vt in data.vertex_types:
- res.add_vertex_type(vt.name, vt.count)
- for key, et in data.edge_types.items():
- adjacency_matrices_neg = []
- for adj_mat in et.adjacency_matrices:
- adj_mat_neg = negative_sample_adj_mat(adj_mat)
- adjacency_matrices_neg.append(adj_mat_neg)
- res.add_edge_type(et.name,
- et.vertex_type_row, et.vertex_type_column,
- adjacency_matrices_neg, et.decoder_factory)
- #new_et = EdgeType(et.name, et.vertex_type_row,
- # et.vertex_type_column, adjacency_matrices_neg,
- # et.decoder_factory, et.total_connectivity)
- #new_edge_types[key] = new_et
- #res = Data(data.vertex_types, new_edge_types)
- return res
-
-
- def merge_data(pos_data: Data, neg_data: Data) -> Data:
- assert isinstance(pos_data, Data)
- assert isinstance(neg_data, Data)
-
- res = PosNegData()
-
- for vt in pos_data.vertex_types:
- res.add_vertex_type(vt.name, vt.count)
-
- for key, pos_et in pos_data.edge_types.items():
- neg_et = neg_data.edge_types[key]
- res.add_edge_type(pos_et.name,
- pos_et.vertex_type_row, pos_et.vertex_type_column,
- pos_et.adjacency_matrices, neg_et.adjacency_matrices,
- pos_et.decoder_factory)
|