diff --git a/src/.DS_Store b/src/.DS_Store new file mode 100644 index 0000000..5008ddf Binary files /dev/null and b/src/.DS_Store differ diff --git a/src/attn_sparsevd.pdf b/src/attn_sparsevd.pdf new file mode 100644 index 0000000..7734abe Binary files /dev/null and b/src/attn_sparsevd.pdf differ diff --git a/src/models/__init__.py b/src/models/__init__.py index 836757f..cb519cb 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -1,4 +1,5 @@ from .lessr import LESSR from .msgifsr import MSGIFSR from .niser import NISER -from .srgnn import SRGNN \ No newline at end of file +from .srgnn import SRGNN +from .niser_ode import NISER_ODE \ No newline at end of file diff --git a/src/models/niser_ode.py b/src/models/niser_ode.py new file mode 100644 index 0000000..cbf2954 --- /dev/null +++ b/src/models/niser_ode.py @@ -0,0 +1,385 @@ +import math + +import torch as th +import torch.nn as nn +import torch.nn.functional as F + +import dgl +import dgl.ops as F +import dgl.function as fn + +from dgl.nn.pytorch import GraphConv, GATConv + +from torchdiffeq import odeint_adjoint, odeint + +from torch.autograd import Variable + +class GraphGRUODE(nn.Module): + + def __init__(self, in_dim, hid_dim, device=th.device('cpu'), gnn='GCNConv', bias=True, **kwargs): + + super(GraphGRUODE, self).__init__() + + self.in_dim = in_dim + self.hid_dim = hid_dim + self.device = device + self.gnn = gnn + self.bias = bias + self.dropout = nn.Dropout(0.1) + + if self.gnn == 'GCNConv': + # self.lin_xx = GCNConv(self.in_dim+self.hid_dim, self.hid_dim, bias=self.bias) + # self.lin_hx = nn.Linear(self.hid_dim, self.in_dim, bias=self.bias) + self.lin_xz = GraphConv(self.in_dim, self.hid_dim, bias=self.bias, allow_zero_in_degree=True) + self.lin_xr = GraphConv(self.in_dim, self.hid_dim, bias=self.bias, allow_zero_in_degree=True) + self.lin_xh = GraphConv(self.in_dim, self.hid_dim, bias=self.bias, allow_zero_in_degree=True) + self.lin_hz = GraphConv(self.hid_dim, self.hid_dim, bias=self.bias, allow_zero_in_degree=True) + self.lin_hr = GraphConv(self.hid_dim, self.hid_dim, bias=self.bias, allow_zero_in_degree=True) + self.lin_hh = GraphConv(self.hid_dim, self.hid_dim, bias=self.bias, allow_zero_in_degree=True) + elif self.gnn == 'GATConv': + self.lin_xz = GATConv(self.in_dim, self.hid_dim, bias=self.bias, num_heads=8, allow_zero_in_degree=True) + self.lin_xr = GATConv(self.in_dim, self.hid_dim, bias=self.bias, num_heads=8, allow_zero_in_degree=True) + self.lin_xh = GATConv(self.in_dim, self.hid_dim, bias=self.bias, num_heads=8, allow_zero_in_degree=True) + self.lin_hz = GATConv(self.hid_dim, self.hid_dim, bias=self.bias, num_heads=8, allow_zero_in_degree=True) + self.lin_hr = GATConv(self.hid_dim, self.hid_dim, bias=self.bias, num_heads=8, allow_zero_in_degree=True) + self.lin_hh = GATConv(self.hid_dim, self.hid_dim, bias=self.bias, num_heads=8, allow_zero_in_degree=True) + elif self.gnn == 'Linear': + self.lin_xx = nn.Linear(self.in_dim, self.hid_dim, bias=self.bias) + self.lin_hx = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + self.lin_xz = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + self.lin_xr = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + self.lin_xh = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + self.lin_hz = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + self.lin_hr = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + self.lin_hh = nn.Linear(self.hid_dim, self.hid_dim, bias=self.bias) + else: + raise NotImplementedError + + self.edge_index = None + self.x = None + + # self.reset_parameters() + + # def reset_parameters(self): + + # self.lin_xz.reset_parameters() + + def set_graph(self, graph: dgl.DGLGraph): + + self.graph = graph + + def set_x(self, x): + self.x = x.to(self.device) + + def forward(self, t, h): + + # x = torch.zeros_like(h).to(self.device) + + # edge_index = self.edge_index_batchs[0] + + # edge_idx = self.graph.filter_edges(lambda edges: edges.data['t'] <= t) + # print(sum(edge_idx.long())) + edge_index = self.graph.edges() + graph = dgl.graph((edge_index[0], edge_index[1]), num_nodes=self.graph.number_of_nodes(), device=self.device) + graph = dgl.remove_self_loop(graph) + graph = dgl.add_reverse_edges(graph) + # graph = dgl.add_self_loop(graph) + # graph = self.graph + # x = self.dropout(self.x) + # h = self.dropout(h) + x = self.x + + + if self.gnn == 'GATConv': + # x = self.lin_xx(torch.cat((self.x.to(self.device), h), dim=1), edge_index).to(self.device) + xr, xz, xh = self.lin_xr(graph, x).max(1)[0], self.lin_xz(graph, x).max(1)[0], self.lin_xh(graph, x).max(1)[0] + r = th.sigmoid(xr + self.lin_hr(graph, h).max(1)[0]) + z = th.sigmoid(xz + self.lin_hz(graph, h).max(1)[0]) + u = th.tanh(xh + self.lin_hh(graph, r * h).max(1)[0]) + elif self.gnn == 'GCNConv': + xr, xz, xh = self.lin_xr(graph, x), self.lin_xz(graph, x), self.lin_xh(graph, x) + r = th.sigmoid(xr + self.lin_hr(graph, h)) + z = th.sigmoid(xz + self.lin_hz(graph, h)) + u = th.tanh(xh + self.lin_hh(graph, r * h)) + else: + # print(h.shape) + h = self.lin_hx(h)+self.lin_xx(x) + # x = self.propagate(edge_index=edge_index, x=h, aggr='mean')-h + xr, xz, xh = self.lin_xr(x), self.lin_xz(x), self.lin_xh(x) + r = th.sigmoid(xr + self.lin_hr(h)) + z = th.sigmoid(xz + self.lin_hz(h)) + u = th.tanh(xh + self.lin_hh(r * h)) + + + dh = (1 - z) * (u - h) + + dh = nn.functional.normalize(dh) + # self.x = self.hx(dh, edge_index) + return dh + +class CDEFunc(nn.Module): + def __init__(self, input_channels, hidden_channels): + ###################### + # input_channels is the number of input channels in the data X. (Determined by the data.) + # hidden_channels is the number of channels for z_t. (Determined by you!) + ###################### + super(CDEFunc, self).__init__() + self.input_channels = input_channels + self.hidden_channels = hidden_channels + + self.linear1 = nn.Linear(hidden_channels, 128) + self.linear2 = nn.Linear(128, input_channels * hidden_channels) + + ###################### + # For most purposes the t argument can probably be ignored; unless you want your CDE to behave differently at + # different times, which would be unusual. But it's there if you need it! + ###################### + + def set_graph(self, graph: dgl.DGLGraph): + + self.graph = graph + + def forward(self, t, z): + # z has shape (batch, hidden_channels) + z = self.linear1(z) + z = z.relu() + z = self.linear2(z) + ###################### + # Easy-to-forget gotcha: Best results tend to be obtained by adding a final tanh nonlinearity. + ###################### + z = z.tanh() + ###################### + # Ignoring the batch dimension, the shape of the output tensor must be a matrix, + # because we need it to represent a linear map from R^input_channels to R^hidden_channels. + ###################### + z = z.view(z.size(0), self.hidden_channels, self.input_channels) + return z + +class GGNNLayer(nn.Module): + def __init__(self, input_dim, output_dim, feat_drop=0.0, activation=None): + super().__init__() + self.dropout = nn.Dropout(feat_drop) + self.gru = nn.GRUCell(2 * output_dim, input_dim) + self.W1 = nn.Linear(input_dim, output_dim, bias=False) + self.W2 = nn.Linear(input_dim, output_dim, bias=False) + self.activation = activation + + def messager(self, edges): + + return {'m': edges.src['ft'] * edges.data['w'].unsqueeze(-1), 'w': edges.data['w']} + + def reducer(self, nodes): + + m = nodes.mailbox['m'] + w = nodes.mailbox['w'] + hn = m.sum(dim=1) / w.sum(dim=1).unsqueeze(-1) + + return {'neigh': hn} + + def forward(self, mg, feat): + with mg.local_scope(): + mg.ndata['ft'] = self.dropout(feat) + if mg.number_of_edges() > 0: + mg.update_all(self.messager, self.reducer) + neigh1 = mg.ndata['neigh'] + mg1 = mg.reverse(copy_edata=True) + mg1.update_all(self.messager, self.reducer) + neigh2 = mg1.ndata['neigh'] + neigh1 = self.W1(neigh1) + neigh2 = self.W2(neigh2) + hn = th.cat((neigh1, neigh2), dim=1) + rst = self.gru(hn, feat) + else: + rst = feat + if self.activation is not None: + rst = self.activation(rst) + return rst + +class GGATLayer(nn.Module): + + def __init__(self, input_dim, output_dim, feat_drop=0.0, activation=None): + super().__init__() + self.dropout = nn.Dropout(feat_drop) + self.gru = nn.GRUCell(2 * input_dim, output_dim) + self.W1 = GATConv(input_dim, output_dim, 8, feat_drop=feat_drop, attn_drop=feat_drop, residual=False, negative_slope=0.1, allow_zero_in_degree=True) + self.W2 = GATConv(input_dim, output_dim, 8, feat_drop=feat_drop, attn_drop=feat_drop, residual=False, negative_slope=0.1, allow_zero_in_degree=True) + self.activation = activation + + def forward(self, mg, feat): + with mg.local_scope(): + mg = dgl.remove_self_loop(mg) + # mg = dgl.add_self_loop(mg) + if mg.number_of_nodes() > 0: + neigh1 = self.W1(mg, feat).max(1)[0] + mg1 = mg.reverse(copy_edata=True) + neigh2 = self.W2(mg1, feat).max(1)[0] + hn = th.cat((neigh1, neigh2), dim=1) + rst = self.gru(hn, feat) + else: + rst = feat + + if self.activation is not None: + rst = self.activation(rst) + + return rst + + +class AttnReadout(nn.Module): + def __init__( + self, + input_dim, + hidden_dim, + output_dim, + batch_norm=True, + feat_drop=0.0, + activation=None, + ): + super().__init__() + self.batch_norm = nn.BatchNorm1d(input_dim) if batch_norm else None + self.feat_drop = nn.Dropout(feat_drop) + self.fc_u = nn.Linear(input_dim, hidden_dim, bias=False) + self.fc_v = nn.Linear(input_dim, hidden_dim, bias=True) + self.fc_e = nn.Linear(hidden_dim, 1, bias=False) + self.fc_out = ( + nn.Linear(input_dim, output_dim, bias=False) + if output_dim != input_dim + else None + ) + self.activation = activation + + def forward(self, g, feat, last_nodes): + if self.batch_norm is not None: + feat = self.batch_norm(feat) + feat = self.feat_drop(feat) + feat_u = self.fc_u(feat) + feat_v = self.fc_v(feat[last_nodes]) + feat_v = dgl.broadcast_nodes(g, feat_v) + e = self.fc_e(th.sigmoid(feat_u + feat_v)) + alpha = F.segment.segment_softmax(g.batch_num_nodes(), e) + feat_norm = feat * alpha + rst = F.segment.segment_reduce(g.batch_num_nodes(), feat_norm, 'sum') + if self.fc_out is not None: + rst = self.fc_out(rst) + if self.activation is not None: + rst = self.activation(rst) + return rst + +class NISER_ODE(nn.Module): + + def __init__(self, num_items, embedding_dim, num_layers, feat_drop=0.0, norm=True, scale=12, solver=None): + super().__init__() + self.num_items = num_items + self.embedding = nn.Embedding(num_items, embedding_dim) + self.register_buffer('indices', th.arange(num_items, dtype=th.long)) + self.embedding_dim = embedding_dim + self.num_layers = num_layers + self.layers = nn.ModuleList() + self.norm = norm + self.scale = scale + self.solver = solver + input_dim = embedding_dim + for i in range(num_layers): + layer = GGNNLayer( + input_dim, + embedding_dim * 2, + feat_drop=feat_drop + ) + self.layers.append(layer) + self.readout = AttnReadout( + input_dim, + embedding_dim, + embedding_dim, + batch_norm=None, + feat_drop=feat_drop, + activation=None, + ) + + self.ODEFunc = GraphGRUODE(self.embedding_dim, self.embedding_dim, device=th.device('cuda:0')) + + # self.initial = GraphConv(self.embedding_dim, 2*self.embedding_dim, allow_zero_in_degree=True) + # self.initial = nn.Linear(self.embedding_dim, self.embedding_dim) + # self.enc_mean = nn.Sequential(nn.Linear(self.embedding_dim, self.embedding_dim), nn.ReLU(), nn.Linear(self.embedding_dim, self.embedding_dim, bias=False)) + # self.enc_var = nn.Sequential(nn.Linear(self.embedding_dim, self.embedding_dim, bias=False), nn.ReLU()) + self.enc_mean = GGNNLayer(input_dim, embedding_dim, feat_drop=feat_drop) + self.enc_var = GGNNLayer(input_dim, embedding_dim, feat_drop=feat_drop) + + self.feat_drop = nn.Dropout(feat_drop) + self.fc_sr = nn.Linear(input_dim + embedding_dim, embedding_dim, bias=False) + + self.reset_parameters() + + def reset_parameters(self): + stdv = 1.0 / math.sqrt(self.embedding_dim) + for weight in self.parameters(): + weight.data.uniform_(-stdv, stdv) + + def _reparameterized_sample(self, mean, std): + eps1 = th.FloatTensor(std.size()).normal_().to(mean.device) + # eps1 = Variable(eps1).to(mean.device) + # return eps1.mul(std).add_(mean) + return mean + eps1 * std + + def forward(self, mg, embeds_ids, times, num_nodes): + + iid = mg.ndata['iid'] + + # print(iid.max(), self.num_items) + feat = self.feat_drop(self.embedding(iid)) + if self.norm: + feat = feat.div(th.norm(feat, p=2, dim=-1, keepdim=True) + 1e-12) + + # feat0 = feat + out = feat + for i, layer in enumerate(self.layers): + out = layer(mg, out) + + feat = out + # feat_mean = self.enc_mean(mg, feat) + # feat_var = nn.functional.relu(self.enc_var(mg, feat)) + # feat = self._reparameterized_sample(feat_mean, feat_var) + # # mgs = dgl.add_reverse_edges(mg) + # # feat0 = feat + # feat0_mean = nn.functional.tanh(feat0[:, :self.embedding_dim]) + # feat0_var = nn.functional.softplus(feat0[:, self.embedding_dim:]) + # feat0 = self._reparameterized_sample(feat0_mean, feat0_var) + + # feat_mean = self.enc_mean(feat) + # feat_var = self.enc_var(feat) + # feat = self._reparameterized_sample(feat_mean, feat_var) + + # feat = self.enc_mean(feat) + + if self.norm: + # feat = feat.div(th.norm(feat, p=2, dim=-1, keepdim=True)) + feat = nn.functional.normalize(feat) + + self.ODEFunc.set_graph(mg) + self.ODEFunc.set_x(feat) + # print(mg.edata) + t_end = mg.edata['t'].max() + + t = th.tensor([0., t_end], device=mg.device) + # print(t) + feat = odeint(self.ODEFunc, feat, t=t, method=self.solver)[-1] # + feat + + last_nodes = mg.filter_nodes(lambda nodes: nodes.data['last'] == 1) + if self.norm: + feat = feat.div(th.norm(feat, p=2, dim=-1, keepdim=True)) + sr_g = self.readout(mg, feat, last_nodes) + sr_l = feat[last_nodes] + sr = th.cat([sr_l, sr_g], dim=1) + sr = self.fc_sr(sr) + if self.norm: + sr = sr.div(th.norm(sr, p=2, dim=-1, keepdim=True) + 1e-12) + target = self.embedding(self.indices) + if self.norm: + target = target.div(th.norm(target, p=2, dim=-1, keepdim=True) + 1e-12) + logits = sr @ target.t() + if self.scale: + logits = th.log(nn.functional.softmax(self.scale * logits, dim=-1)) + else: + logits = th.log(nn.functional.softmax(logits, dim=-1)) + return logits# , 0 + + + \ No newline at end of file diff --git a/src/preprocess.py b/src/preprocess.py new file mode 100644 index 0000000..221c177 --- /dev/null +++ b/src/preprocess.py @@ -0,0 +1,52 @@ +from pathlib import Path +import argparse + + +parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + +optional = parser._action_groups.pop() +required = parser.add_argument_group('required arguments') +required.add_argument( + '-d', + '--dataset', + choices=['diginetica', 'gowalla', 'lastfm'], + required=True, + help='the dataset name', +) +required.add_argument( + '-f', + '--filepath', + required=True, + help='the file for the dataset, i.e., "train-item-views.csv" for diginetica, ' + '"loc-gowalla_totalCheckins.txt" for gowalla, ' + '"userid-timestamp-artid-artname-traid-traname.tsv" for lastfm', +) +optional.add_argument( + '-t', + '--dataset-dir', + default='datasets/{dataset}', + help='the folder to save the preprocessed dataset', +) +parser._action_groups.append(optional) +args = parser.parse_args() + +dataset_dir = Path(args.dataset_dir.format(dataset=args.dataset)) + +if args.dataset == 'diginetica': + from utils.data.preprocess import preprocess_diginetica + + preprocess_diginetica(dataset_dir, args.filepath) +else: + from pandas import Timedelta + from utils.data.preprocess import preprocess_gowalla_lastfm + + csv_file = args.filepath + if args.dataset == 'gowalla': + usecols = [0, 1, 4] + interval = Timedelta(days=1) + n = 30000 + else: + usecols = [0, 1, 2] + interval = Timedelta(hours=8) + n = 40000 + preprocess_gowalla_lastfm(dataset_dir, csv_file, usecols, interval, n) \ No newline at end of file diff --git a/src/scripts/main_niser_ode.py b/src/scripts/main_niser_ode.py new file mode 100644 index 0000000..91d9230 --- /dev/null +++ b/src/scripts/main_niser_ode.py @@ -0,0 +1,160 @@ +import argparse +import sys +import torch +import random +import numpy as np +import os +import wandb + + +def seed_torch(seed=42): + seed = int(seed) + random.seed(seed) + os.environ['PYTHONHASHSEED'] = str(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.enabled = True + +seed_torch(123) + +def get_freer_gpu(): + os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp') + memory_available = [int(x.split()[2]) for x in open('tmp', 'r').readlines()] + # memory_available = memory_available[1:6] + if len(memory_available) == 0: + return -1 + return int(np.argmax(memory_available)) + +os.environ["CUDA_VISIBLE_DEVICES"] = str(get_freer_gpu()) + +sys.path.append('..') +sys.path.append('../..') + +parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) +parser.add_argument( + '--dataset-dir', default='datasets/sample', help='the dataset directory' +) +parser.add_argument( + '--solver', default='rk4', help='The neural ordinary equation solver.' +) +parser.add_argument('--embedding-dim', type=int, default=256, help='the embedding size') +parser.add_argument('--num-layers', type=int, default=1, help='the number of layers') +parser.add_argument( + '--feat-drop', type=float, default=0.1, help='the dropout ratio for features' +) +parser.add_argument('--lr', type=float, default=1e-3, help='the learning rate') +parser.add_argument( + '--batch-size', type=int, default=512, help='the batch size for training' +) +parser.add_argument( + '--epochs', type=int, default=30, help='the number of training epochs' +) +parser.add_argument( + '--weight-decay', + type=float, + default=1e-4, + help='the parameter for L2 regularization', +) +parser.add_argument( + '--patience', + type=int, + default=2, + help='the number of epochs that the performance does not improves after which the training stops', +) +parser.add_argument( + '--num-workers', + type=int, + default=10, + help='the number of processes to load the input graphs', +) +parser.add_argument( + '--valid-split', + type=float, + default=None, + help='the fraction for the validation set', +) +parser.add_argument( + '--log-interval', + type=int, + default=100, + help='print the loss after this number of iterations', +) +args = parser.parse_args() +print(args) + +wandb.init(config=vars(args)) + + +from pathlib import Path +import torch as th +from torch.utils.data import DataLoader, SequentialSampler +from src.utils.data.dataset import read_dataset, AugmentedDataset +from src.utils.data.collate import ( + seq_to_temporal_session_graph, + collate_fn_factory_temporal, +) +from src.utils.train import TrainRunner +from src.models import NISER_ODE + +dataset_dir = Path(args.dataset_dir) + +print('reading dataset') +train_sessions, test_sessions, train_timestamps, test_timestamps, num_items = read_dataset(dataset_dir) + +if args.valid_split is not None: + num_valid = int(len(train_sessions) * args.valid_split) + test_sessions = train_sessions[-num_valid:] + train_sessions = train_sessions[:-num_valid] + test_timestamps = train_timestamps[-num_valid:] + train_timestamps = train_timestamps[:-num_valid] + +dataset = args.dataset_dir.strip().split("/")[-1] + +train_set = AugmentedDataset(dataset, train_sessions, train_timestamps) +test_set = AugmentedDataset(dataset, test_sessions, test_timestamps) + +collate_fn = collate_fn_factory_temporal(seq_to_temporal_session_graph) + +train_loader = DataLoader( + train_set, + batch_size=args.batch_size, + shuffle=True, + # drop_last=True, + num_workers=args.num_workers, + collate_fn=collate_fn, + pin_memory=True, + # sampler=SequentialSampler(train_set) +) + +test_loader = DataLoader( + test_set, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + collate_fn=collate_fn, +) + +model = NISER_ODE(num_items, args.embedding_dim, args.num_layers, feat_drop=args.feat_drop, solver=args.solver) +device = th.device('cuda:0' if th.cuda.is_available() else 'cpu') +model = model.to(device) +print(model) + +runner = TrainRunner( + args.dataset_dir, + model, + train_loader, + test_loader, + device=device, + lr=args.lr, + weight_decay=args.weight_decay, + patience=args.patience, +) + +print('start training') +mrr10, mrr20, hit10, hit20 = runner.train(args.epochs, args.log_interval) +print('MRR@20\tHR@20') +print(f'{mrr10 * 100:.3f}%\t{mrr20 * 100:.3f}%\t{hit10 * 100:.3f}%\t{hit20 * 100:.3f}%') diff --git a/src/sweep.yaml b/src/sweep.yaml new file mode 100644 index 0000000..54f15fc --- /dev/null +++ b/src/sweep.yaml @@ -0,0 +1,9 @@ +project: gng-ode-solvers +program: scripts/main_niser_ode.py +method: grid + +parameters: + solver: + values: ["euler", "rk4", "implicit_adams"] + dataset-dir: + values: ["../datasets/gowalla", "../datasets/tmall", "../datasets/nowplaying"] \ No newline at end of file diff --git a/src/tmall_hit_solver.pdf b/src/tmall_hit_solver.pdf new file mode 100644 index 0000000..74de46a Binary files /dev/null and b/src/tmall_hit_solver.pdf differ diff --git a/src/utils/__pycache__/train.cpython-38.pyc b/src/utils/__pycache__/train.cpython-38.pyc index fd7111a..415b8f5 100644 Binary files a/src/utils/__pycache__/train.cpython-38.pyc and b/src/utils/__pycache__/train.cpython-38.pyc differ diff --git a/src/utils/data/__pycache__/collate.cpython-38.pyc b/src/utils/data/__pycache__/collate.cpython-38.pyc index adce854..d1fc3d6 100644 Binary files a/src/utils/data/__pycache__/collate.cpython-38.pyc and b/src/utils/data/__pycache__/collate.cpython-38.pyc differ diff --git a/src/utils/data/__pycache__/collate.cpython-39.pyc b/src/utils/data/__pycache__/collate.cpython-39.pyc index a7596a4..c980daa 100644 Binary files a/src/utils/data/__pycache__/collate.cpython-39.pyc and b/src/utils/data/__pycache__/collate.cpython-39.pyc differ diff --git a/src/utils/data/__pycache__/dataset.cpython-38.pyc b/src/utils/data/__pycache__/dataset.cpython-38.pyc index df4e821..c19cb1a 100644 Binary files a/src/utils/data/__pycache__/dataset.cpython-38.pyc and b/src/utils/data/__pycache__/dataset.cpython-38.pyc differ diff --git a/src/utils/data/__pycache__/dataset.cpython-39.pyc b/src/utils/data/__pycache__/dataset.cpython-39.pyc index 73770ad..2aeae27 100644 Binary files a/src/utils/data/__pycache__/dataset.cpython-39.pyc and b/src/utils/data/__pycache__/dataset.cpython-39.pyc differ diff --git a/src/utils/data/__pycache__/preprocess.cpython-38.pyc b/src/utils/data/__pycache__/preprocess.cpython-38.pyc index d791b54..d2715d6 100644 Binary files a/src/utils/data/__pycache__/preprocess.cpython-38.pyc and b/src/utils/data/__pycache__/preprocess.cpython-38.pyc differ diff --git a/src/utils/data/collate.py b/src/utils/data/collate.py index db55c7d..6c17176 100644 --- a/src/utils/data/collate.py +++ b/src/utils/data/collate.py @@ -1,6 +1,7 @@ from collections import Counter import numpy as np import torch as th +import torch.nn.functional as F import dgl import pickle import numba @@ -84,6 +85,42 @@ def seq_to_session_graph(seq): return g +def seq_to_temporal_session_graph(seq, times): + items, indices = np.unique(seq, return_index=True) + iid2nid = {iid: i for i, iid in enumerate(items)} + num_nodes = len(items) + + seq_nid = [iid2nid[iid] for iid in seq] + counter = Counter( + [(seq_nid[i], seq_nid[i+1]) for i in range(len(seq)-1)] + ) + edges = counter.keys() + if len(edges) > 0: + src, dst = zip(*edges) + weight = th.tensor(list(counter.values())) + else: + src, dst = [0], [0] + weight = th.ones(1).long() + + g = dgl.graph((src, dst), num_nodes=num_nodes) + + g.edata['w'] = weight + # print(len(times), g.number_of_nodes()) + g.ndata['t'] = th.tensor(times)[indices] + + # print(g.edata, times, g.number_of_edges(), g.number_of_nodes()) + if g.number_of_edges() == 1 and g.number_of_nodes() == 1: + g.edata['t'] = th.tensor(times)[0].unsqueeze(-1) + else: + g.edata['t'] = th.tensor(times)[-g.number_of_edges():] + + # print(g.edata) + + g.ndata['iid'] = th.from_numpy(items) + label_last(g, iid2nid[seq[-1]]) + + return g + def seq_to_ccs_graph(seq, order=1, coaDict=None): order1 = order @@ -229,6 +266,24 @@ def collate_fn(samples): return collate_fn +def collate_fn_factory_temporal(*seq_to_graph_fns): + def collate_fn(samples): + seqs, times, labels = zip(*samples) + inputs = [] + for seq_to_graph in seq_to_graph_fns: + graphs = list(map(seq_to_graph, seqs, times)) + num_nodes = th.tensor([graph.number_of_nodes() for graph in graphs], dtype=th.long) + max_num = max(num_nodes) + embeds_id = th.vstack([F.pad(graph.ndata['iid'], (0, max_num-len(graph.ndata['iid'])), value=graph.ndata['iid'][-1]) for graph in graphs]) + times = th.vstack([F.pad(graph.ndata['t'], (0, max_num-len(graph.ndata['iid'])), value=graph.ndata['iid'][-1]) for graph in graphs]) + bg = dgl.batch(graphs) + inputs.append(bg) + labels = th.LongTensor(labels) + # print(inputs[0].edata) + return inputs, labels, embeds_id, times, num_nodes + + return collate_fn + def collate_fn_factory_ccs(seq_to_graph_fns, order): def collate_fn(samples): seqs, labels = zip(*samples) diff --git a/src/utils/data/dataset.py b/src/utils/data/dataset.py index 5e3e38c..24942e7 100644 --- a/src/utils/data/dataset.py +++ b/src/utils/data/dataset.py @@ -1,6 +1,8 @@ import itertools +from os import read import numpy as np import pandas as pd +import pickle as pkl def create_index(sessions): @@ -18,17 +20,36 @@ def read_sessions(filepath): sessions = sessions.apply(lambda x: list(map(int, x.split(',')))).values return sessions +def read_timestamps(filepath): + sessions = pd.read_csv(filepath, sep='\t', header=None, squeeze=True) + sessions = sessions.apply(lambda x: list(map(float, x.split(',')))).values + return sessions def read_dataset(dataset_dir): - train_sessions = read_sessions(dataset_dir / 'train.txt') - test_sessions = read_sessions(dataset_dir / 'test.txt') + dataset = dataset_dir.name.strip().split('/')[-1] + if dataset in ["gowalla"]: + train_sessions = read_sessions(dataset_dir / 'train.txt') + test_sessions = read_sessions(dataset_dir / 'test.txt') + train_timestamp = read_timestamps(dataset_dir / 'train_timestamp.txt') + test_timestamp = read_timestamps(dataset_dir / 'test_timestamp.txt') + elif dataset in ["tmall", "nowplaying"]: + train_dict = pkl.load(open(dataset_dir / "train.txt", "rb")) + test_dict = pkl.load(open(dataset_dir / "test.txt", "rb")) + train_sessions = train_dict[0] + test_sessions = test_dict[0] + train_timestamp = train_dict[1] + test_timestamp = test_dict[1] + with open(dataset_dir / 'num_items.txt', 'r') as f: - num_items = int(f.readline()) - return train_sessions, test_sessions, num_items + num_items = int(f.readline()) + return train_sessions, test_sessions, train_timestamp, test_timestamp, num_items + class AugmentedDataset: - def __init__(self, sessions, sort_by_length=False): - self.sessions = sessions + def __init__(self, dataset, sessions, timestamps, sort_by_length=False): + self.dataset = dataset + self.sessions = sessions + self.timestamps = timestamps # self.graphs = graphs index = create_index(sessions) # columns: sessionId, labelIndex @@ -41,10 +62,24 @@ def __init__(self, sessions, sort_by_length=False): def __getitem__(self, idx): #print(idx) sid, lidx = self.index[idx] - seq = self.sessions[sid][:lidx] - label = self.sessions[sid][lidx] + seq = self.sessions[sid][:lidx] + label = self.sessions[sid][lidx] + times = self.timestamps[sid][:lidx]# - self.sessions[sid][0] + temp0 = times[0] + # temp = times[-1] + + if self.dataset in ["tmall"]: + scale = 100 + elif self.dataset in ["nowplaying"]: + scale = 100000 + else: + scale = 1000000 + + times = [(t - temp0) / scale for t in times] + + # print(times) - return seq, label #,seq + return seq, times, label #,seq def __len__(self): return len(self.index) diff --git a/src/utils/data/preprocess.py b/src/utils/data/preprocess.py index fcf9efd..fba6394 100644 --- a/src/utils/data/preprocess.py +++ b/src/utils/data/preprocess.py @@ -1,6 +1,6 @@ import pandas as pd import numpy as np - +import time def get_session_id(df, interval): df_prev = df.shift() @@ -103,9 +103,17 @@ def train_test_split(df, test_split=0.2): def save_sessions(df, filepath): df = reorder_sessions_by_endtime(df) - sessions = df.groupby('sessionId').itemId.apply(lambda x: ','.join(map(str, x))) + sessions = df.groupby('sessionId')# . + sessions = sessions.itemId.apply(lambda x: ','.join(map(str, x))) sessions.to_csv(filepath, sep='\t', header=False, index=False) - + # sessions_timestamp = sessions.timestamp.apply(lambda x: ','.join(map(str, x))) + +def save_sessions_timestamp(df, filepath): + df = reorder_sessions_by_endtime(df) + df['timestamp'] = df['timestamp'].apply(lambda x: time.mktime(x.timetuple())) + sessions = df.groupby('sessionId')# . + sessions = sessions.timestamp.apply(lambda x: ','.join(map(str, x))) + sessions.to_csv(filepath, sep='\t', header=False, index=False) def save_dataset(dataset_dir, df_train, df_test): # filter items in test but not in train @@ -117,15 +125,17 @@ def save_dataset(dataset_dir, df_train, df_test): # update itemId train_itemId_new, uniques = pd.factorize(df_train.itemId) - df_train = df_train.assign(itemId=train_itemId_new) - oid2nid = {oid: i for i, oid in enumerate(uniques)} - test_itemId_new = df_test.itemId.map(oid2nid) - df_test = df_test.assign(itemId=test_itemId_new) + df_train = df_train.assign(itemId=train_itemId_new) + oid2nid = {oid: i for i, oid in enumerate(uniques)} + test_itemId_new = df_test.itemId.map(oid2nid) + df_test = df_test.assign(itemId=test_itemId_new) print(f'saving dataset to {dataset_dir}') dataset_dir.mkdir(parents=True, exist_ok=True) save_sessions(df_train, dataset_dir / 'train.txt') - save_sessions(df_test, dataset_dir / 'test.txt') + save_sessions(df_test, dataset_dir / 'test.txt') + save_sessions_timestamp(df_train, dataset_dir / 'train_timestamp.txt') + save_sessions_timestamp(df_test, dataset_dir / 'test_timestamp.txt') num_items = len(uniques) with open(dataset_dir / 'num_items.txt', 'w') as f: f.write(str(num_items)) @@ -154,7 +164,6 @@ def preprocess_diginetica(dataset_dir, csv_file): df_train, df_test = split_by_time(df, pd.Timedelta(days=7)) save_dataset(dataset_dir, df_train, df_test) - def preprocess_gowalla_lastfm(dataset_dir, csv_file, usecols, interval, n): print(f'reading {csv_file}...') df = pd.read_csv( diff --git a/src/utils/train.py b/src/utils/train.py index 4de090a..a65b84d 100644 --- a/src/utils/train.py +++ b/src/utils/train.py @@ -24,12 +24,15 @@ def fix_weight_decay(model): def prepare_batch(batch, device): - inputs, labels = batch + inputs, labels, embeds_ids, times, num_nodes = batch # inputs, labels = batch inputs_gpu = [x.to(device) for x in inputs] labels_gpu = labels.to(device) + embeds_ids = embeds_ids.to(device) + times = times.to(device) + num_nodes = num_nodes.to(device) - return inputs_gpu, labels_gpu + return inputs_gpu, labels_gpu, embeds_ids, times, num_nodes # return inputs_gpu, 0, labels_gpu, 0 @@ -41,8 +44,8 @@ def evaluate(model, data_loader, device, cutoff=20): with th.no_grad(): for batch in data_loader: - inputs, labels = prepare_batch(batch, device) - logits = model(*inputs) + inputs, labels, embeds_ids, times, num_nodes = prepare_batch(batch, device) + logits = model(*inputs, embeds_ids, times, num_nodes) batch_size = logits.size(0) num_samples += batch_size @@ -82,8 +85,8 @@ def __init__( self.patience = patience def train(self, epochs, log_interval=100): - max_mrr = 0 - max_hit = 0 + max_mrr10 = max_mrr20 = 0 + max_hit10 = max_hit20 = 0 bad_counter = 0 t = time.time() mean_loss = 0 @@ -92,9 +95,9 @@ def train(self, epochs, log_interval=100): for epoch in tqdm(range(epochs)): self.model.train() for batch in self.train_loader: - inputs, labels = prepare_batch(batch, self.device) + inputs, labels, embeds_ids, times, num_nodes = prepare_batch(batch, self.device) self.optimizer.zero_grad() - scores = self.model(*inputs) + scores = self.model(*inputs, embeds_ids, times, num_nodes) assert not th.isnan(scores).any() loss = nn.functional.nll_loss(scores, labels) loss.backward() @@ -109,19 +112,23 @@ def train(self, epochs, log_interval=100): self.batch += 1 self.scheduler.step() - mrr, hit = evaluate(self.model, self.test_loader, self.device) + mrr10, hit10 = evaluate(self.model, self.test_loader, self.device, cutoff=10) + mrr20, hit20 = evaluate(self.model, self.test_loader, self.device) - # wandb.log({"hit": hit, "mrr": mrr}) + wandb.log({"hit@20": hit20, "mrr@20": mrr20}) - print(f'Epoch {self.epoch}: MRR = {mrr * 100:.3f}%, Hit = {hit * 100:.3f}%') + print(f'Epoch {self.epoch}: MRR@10 = {mrr10 * 100:.3f}%, Hit@10 = {hit10 * 100:.3f}%, MRR@20 = {mrr20 * 100:.3f}%, Hit@20 = {hit20 * 100:.3f}%') - if mrr < max_mrr and hit < max_hit: + if mrr20 < max_mrr20 and hit20 < max_hit20: bad_counter += 1 if bad_counter == self.patience: break else: bad_counter = 0 - max_mrr = max(max_mrr, mrr) - max_hit = max(max_hit, hit) + max_mrr10 = max(max_mrr10, mrr10) + max_hit10 = max(max_hit10, hit10) + max_mrr20 = max(max_mrr20, mrr20) + max_hit20 = max(max_hit20, hit20) + self.epoch += 1 - return max_mrr, max_hit + return max_mrr10, max_mrr20, max_hit10, max_hit20