五、Jbeil utils
1.utils
这个文件实现了几个核心功能:网络结构的定义、早停机制、负边采样器、邻居查找器和邻居采样器。
- 定义了一些神经网络层,用于处理输入特征并进行非线性变换。
- 实现了一个早停机制,用于防止过拟合。
- 定义了一个随机边采样器,用于从给定的节点列表中随机采样边。
- 构建了一个邻居查找器,用于查找给定节点的邻居并进行邻居采样。
MergeLayer 类
class MergeLayer(torch.nn.Module):
def __init__(self, dim1, dim2, dim3, dim4):
super().__init__()
self.fc1 = torch.nn.Linear(dim1 + dim2, dim3)
self.fc2 = torch.nn.Linear(dim3, dim4)
self.act = torch.nn.ReLU()
torch.nn.init.xavier_normal_(self.fc1.weight)
torch.nn.init.xavier_normal_(self.fc2.weight)
def forward(self, x1, x2):
x = torch.cat([x1, x2], dim=1)
h = self.act(self.fc1(x))
return self.fc2(h)
这个类定义了一个简单的神经网络层,它将两个输入向量拼接起来,然后通过两层全连接层和一个ReLU激活函数。主要用于将两个向量合并并进行非线性变换。
MLP 类
class MLP(torch.nn.Module):
def __init__(self, dim, drop=0.3):
super().__init__()
self.fc_1 = torch.nn.Linear(dim, 80)
self.fc_2 = torch.nn.Linear(80, 10)
self.fc_3 = torch.nn.Linear(10, 1)
self.act = torch.nn.ReLU()
self.dropout = torch.nn.Dropout(p=drop, inplace=False)
def forward(self, x):
x = self.act(self.fc_1(x))
x = self.dropout(x)
x = self.act(self.fc_2(x))
x = self.dropout(x)
return self.fc_3(x).squeeze(dim=1)
这个类定义了一个多层感知器(MLP),包括三层全连接层和ReLU激活函数,以及Dropout层用于防止过拟合。
EarlyStopMonitor 类
class EarlyStopMonitor(object):
def __init__(self, max_round=3, higher_better=True, tolerance=1e-10):
self.max_round = max_round
self.num_round = 0
self.epoch_count = 0
self.best_epoch = 0
self.last_best = None
self.higher_better = higher_better
self.tolerance = tolerance
def early_stop_check(self, curr_val):
if not self.higher_better:
curr_val *= -1
if self.last_best is None:
self.last_best = curr_val
elif (curr_val - self.last_best) / np.abs(self.last_best) > self.tolerance:
self.last_best = curr_val
self.num_round = 0
self.best_epoch = self.epoch_count
else:
self.num_round += 1
self.epoch_count += 1
return self.num_round >= self.max_round
该类实现了一个早停机制,用于在训练过程中监控验证集性能,并在性能不再提升时提前停止训练,以防止过拟合。
RandEdgeSampler 类
class RandEdgeSampler(object):
def __init__(self, src_list, dst_list, seed=None):
self.seed = seed
self.src_list = np.unique(src_list)
self.dst_list = np.unique(dst_list)
if seed is not None:
self.random_state = np.random.RandomState(self.seed)
def sample(self, size):
if self.seed is None:
src_index = np.random.randint(0, len(self.src_list), size)
dst_index = np.random.randint(0, len(self.dst_list), size)
else:
src_index = self.random_state.randint(0, len(self.src_list), size)
dst_index = self.random_state.randint(0, len(self.dst_list), size)
return self.src_list[src_index], self.dst_list[dst_index]
def reset_random_state(self):
self.random_state = np.random.RandomState(self.seed)
该类用于从给定的源和目标节点列表中随机采样一组边。它可以使用一个固定的种子来保证采样的可重复性。
get_neighbor_finder 函数
def get_neighbor_finder(data, uniform, max_node_idx=None):
max_node_idx = max(data.sources.max(), data.destinations.max()) if max_node_idx is None else max_node_idx
adj_list = [[] for _ in range(max_node_idx + 1)]
for source, destination, edge_idx, timestamp in zip(data.sources, data.destinations,
data.edge_idxs,
data.timestamps):
adj_list[source].append((destination, edge_idx, timestamp))
adj_list[destination].append((source, edge_idx, timestamp))
return NeighborFinder(adj_list, uniform=uniform)
该函数创建一个邻居查找器对象(NeighborFinder),用于查找给定节点的邻居。它根据输入的数据构建一个邻接列表,其中每个节点都记录了与其相连的节点、边索引和时间戳。
NeighborFinder 类
class NeighborFinder:
def __init__(self, adj_list, uniform=False, seed=None):
self.node_to_neighbors = []
self.node_to_edge_idxs = []
self.node_to_edge_timestamps = []
for neighbors in adj_list:
sorted_neighhbors = sorted(neighbors, key=lambda x: x[2])
self.node_to_neighbors.append(np.array([x[0] for x in sorted_neighhbors]))
self.node_to_edge_idxs.append(np.array([x[1] for x in sorted_neighhbors]))
self.node_to_edge_timestamps.append(np.array([x[2] for x in sorted_neighhbors]))
self.uniform = uniform
if seed is not None:
self.seed = seed
self.random_state = np.random.RandomState(self.seed)
def find_before(self, src_idx, cut_time):
i = np.searchsorted(self.node_to_edge_timestamps[src_idx], cut_time)
return self.node_to_neighbors[src_idx][:i], self.node_to_edge_idxs[src_idx][:i], self.node_to_edge_timestamps[src_idx][:i]
def get_temporal_neighbor(self, source_nodes, timestamps, n_neighbors=20):
assert (len(source_nodes) == len(timestamps))
tmp_n_neighbors = n_neighbors if n_neighbors > 0 else 1
neighbors = np.zeros((len(source_nodes), tmp_n_neighbors)).astype(np.int32)
edge_times = np.zeros((len(source_nodes), tmp_n_neighbors)).astype(np.float32)
edge_idxs = np.zeros((len(source_nodes), tmp_n_neighbors)).astype(np.int32)
for i, (source_node, timestamp) in enumerate(zip(source_nodes, timestamps)):
source_neighbors, source_edge_idxs, source_edge_times = self.find_before(source_node, timestamp)
if len(source_neighbors) > 0 and n_neighbors > 0:
if self.uniform:
sampled_idx = np.random.randint(0, len(source_neighbors), n_neighbors)
neighbors[i, :] = source_neighbors[sampled_idx]
edge_times[i, :] = source_edge_times[sampled_idx]
edge_idxs[i, :] = source_edge_idxs[sampled_idx]
pos = edge_times[i, :].argsort()
neighbors[i, :] = neighbors[i, :][pos]
edge_times[i, :] = edge_times[i, :][pos]
edge_idxs[i, :] = edge_idxs[i, :][pos]
else:
source_edge_times = source_edge_times[-n_neighbors:]
source_neighbors = source_neighbors[-n_neighbors:]
source_edge_idxs = source_edge_idxs[-n_neighbors:]
neighbors[i, n_neighbors - len(source_neighbors):] = source_neighbors
edge_times[i, n_neighbors - len(source_edge_times):] = source_edge_times
edge_idxs[i, n_neighbors - len(source_edge_idxs):] = source_edge_idxs
return neighbors, edge_idxs, edge_times
该类实现了一个邻居查找器,用于查找给定节点的邻居。它可以基于时间戳查找邻居,并根据不同的采样策略(如统一采样)进行采样。主要方法包括:
find_before
: 查找给定节点在特定时间点之前的所有邻居。get_temporal_neighbor
: 获取一组节点在给定时间戳之前的临时邻居。
2.preprocess_data.py
这个preprocess_data
文件的主要功能是预处理输入数据,并生成适合于时间图神经网络(TGN)使用的格式化数据。主要步骤包括数据的读取、特征提取、索引重排和保存。
trunc(values, decs=0)
: 用于截断浮点数到指定的小数位数。preprocess(data_name)
: 读取数据文件并提取出用户、项目、时间戳、标签和特征等信息。reindex(df, bipartite=True)
: 重新排列数据框的索引,并处理二部图(bipartite)和非二部图的情况。run(data_name, bipartite=True)
: 调用上述函数来处理输入数据,并将处理后的数据保存到指定的路径中。
函数详细解释
trunc(values, decs=0)
def trunc(values, decs=0):
return np.trunc(values*10**decs)/(10**decs)
这个函数用于截断浮点数到指定的小数位数。参数values
是需要截断的浮点数,decs
是小数位数。
preprocess(data_name)
def preprocess(data_name):
u_list, i_list, ts_list, label_list = [], [], [], []
feat_l = []
idx_list = []
with open(data_name) as f:
s = next(f)
for idx, line in enumerate(f):
e = line.strip().split(',')
u = int(e[0])
i = int(e[1])
ts = float(e[2])
label = float(e[3])
feat = np.array([float(x) for x in e[4:]])
u_list.append(u)
i_list.append(i)
ts_list.append(ts)
label_list.append(label)
idx_list.append(idx)
feat_l.append(feat)
return pd.DataFrame({'u': u_list,
'i': i_list,
'ts': ts_list,
'label': label_list,
'idx': idx_list}), np.array(feat_l)
这个函数读取数据文件,并将数据解析成用户(u
)、项目(i
)、时间戳(ts
)、标签(label
)和特征(feat
)的列表。然后将这些列表转化为一个DataFrame和一个特征矩阵。
reindex(df, bipartite=True)
def reindex(df, bipartite=True):
new_df = df.copy()
if bipartite:
assert (df.u.max() - df.u.min() + 1 == len(df.u.unique()))
assert (df.i.max() - df.i.min() + 1 == len(df.i.unique()))
upper_u = df.u.max() + 1
new_i = df.i + upper_u
new_df.i = new_i
new_df.u += 1
new_df.i += 1
new_df.idx += 1
else:
new_df.u += 1
new_df.i += 1
new_df.idx += 1
return new_df
这个函数用于重新排列数据框的索引。如果是二部图(bipartite=True
),它会将项目索引偏移,使用户和项目的索引不重叠。对于非二部图,只是简单地将索引加1。
run(data_name, bipartite=True)
(不重要)
def run(data_name, bipartite=True):
Path("data/").mkdir(parents=True, exist_ok=True)
PATH = './data/{}.csv'.format(data_name)
OUT_DF = './data/ml_{}.csv'.format(data_name)
OUT_FEAT = './data/ml_{}.npy'.format(data_name)
OUT_NODE_FEAT = './data/ml_{}_node.npy'.format(data_name)
df, feats = preprocess(PATH)
new_df = reindex(df, bipartite)
empty = np.zeros(feats.shape[1])[np.newaxis, :]
feat = np.vstack([empty, feats])
max_idx = max(new_df.u.max(), new_df.i.max())
rand_feat = np.zeros((max_idx + 1, 10))
new_df.to_csv(OUT_DF)
np.save(OUT_FEAT, feat)
np.save(OUT_NODE_FEAT, rand_feat)
这个函数是预处理数据的主要入口。它创建必要的文件夹,调用preprocess
函数读取和处理数据,然后调用reindex
函数重新排列索引,最后将处理后的数据保存为CSV和Numpy数组文件。
脚本入口
parser = argparse.ArgumentParser('Interface for TGN data preprocessing')
parser.add_argument('--data', type=str, help='Dataset name (eg. auth or pivoting)')
parser.add_argument('--bipartite', action='store_true', help='Whether the graph is bipartite')
args = parser.parse_args()
run(args.data, bipartite=args.bipartite)
这个部分是脚本的入口,使用argparse
解析命令行参数,并调用run
函数进行数据预处理。
函数和类之间的依赖关系
run
函数调用preprocess
函数来读取和初步处理数据。run
函数调用reindex
函数来重新排列索引。run
函数保存处理后的数据到文件中。