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..c8eb215 Binary files /dev/null and b/src/attn_sparsevd.pdf differ diff --git a/src/diginetica_srgnn_sparsevd.pdf b/src/diginetica_srgnn_sparsevd.pdf new file mode 100644 index 0000000..ed9c731 Binary files /dev/null and b/src/diginetica_srgnn_sparsevd.pdf differ diff --git a/src/models/gru.py b/src/models/gru.py new file mode 100644 index 0000000..6303773 --- /dev/null +++ b/src/models/gru.py @@ -0,0 +1,51 @@ +import torch.nn as nn +import torch.nn.functional as F +from .sparsevd import LinearSVDO +import math + +class GRUCell(nn.Module): + + """ + An implementation of GRUCell. + + """ + + def __init__(self, input_size, hidden_size, bias=True, threshold=8.0, name=None): + super(GRUCell, self).__init__() + self.name = name + self.input_size = input_size + self.hidden_size = hidden_size + self.bias = bias + self.x2h = LinearSVDO(input_size, 3 * hidden_size, bias=bias, threshold=threshold, name=name+"_x2h") + self.h2h = LinearSVDO(hidden_size, 3 * hidden_size, bias=bias, threshold=threshold, name=name+"_h2h") + self.reset_parameters() + + + + def reset_parameters(self): + std = 1.0 / math.sqrt(self.hidden_size) + for w in self.parameters(): + w.data.uniform_(-std, std) + + def forward(self, x, hidden): + + x = x.view(-1, x.size(1)) + + gate_x = self.x2h(x) + gate_h = self.h2h(hidden) + + gate_x = gate_x.squeeze() + gate_h = gate_h.squeeze() + + i_r, i_i, i_n = gate_x.chunk(3, 1) + h_r, h_i, h_n = gate_h.chunk(3, 1) + + + resetgate = F.sigmoid(i_r + h_r) + inputgate = F.sigmoid(i_i + h_i) + newgate = F.tanh(i_n + (resetgate * h_n)) + + hy = newgate + inputgate * (hidden - newgate) + + + return hy \ No newline at end of file diff --git a/src/models/msgifsr.py b/src/models/msgifsr.py index d93f207..2f8753b 100644 --- a/src/models/msgifsr.py +++ b/src/models/msgifsr.py @@ -270,8 +270,7 @@ def forward(self, mg): sr = th.cat([sr_l, sr_g], dim=-1)# .view(sr_l.size(0), -1) sr = th.cat([self.fc_sr[i](sr).unsqueeze(1) for i, sr in enumerate(th.unbind(sr, dim=1))], dim=1) if self.norm: - sr = nn.functional.normalize(sr, dim=-1) - + sr = nn.functional.normalize(sr, dim=-1) target = self.embeddings(self.indices) diff --git a/src/models/sparsevd.py b/src/models/sparsevd.py new file mode 100644 index 0000000..83136f1 --- /dev/null +++ b/src/models/sparsevd.py @@ -0,0 +1,46 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Parameter + +class LinearSVDO(nn.Module): + def __init__(self, in_features, out_features, threshold, bias=True, name=None): + super(LinearSVDO, self).__init__() + self.name = name + self.in_features = in_features + self.out_features = out_features + self.threshold = threshold + + self.hasbias = bias + + self.W = Parameter(torch.Tensor(out_features, in_features)) + self.log_sigma = Parameter(torch.Tensor(out_features, in_features)) + self.bias = Parameter(torch.Tensor(1, out_features)) + + self.reset_parameters() + + def reset_parameters(self): + self.bias.data.zero_() + self.W.data.normal_(0, 0.02) + # nn.init.xavier_uniform_(self.W) + self.log_sigma.data.fill_(-5) + + def forward(self, x): + self.log_alpha = self.log_sigma * 2.0 - 2.0 * torch.log(1e-16 + torch.abs(self.W)) + self.log_alpha = torch.clamp(self.log_alpha, -10, 10) + + if self.training: + lrt_mean = F.linear(x, self.W) + self.bias + lrt_std = torch.sqrt(F.linear(x * x, torch.exp(self.log_sigma * 2.0)) + 1e-8) + eps = lrt_std.data.new(lrt_std.size()).normal_() + return lrt_mean + lrt_std * eps + + return F.linear(x, self.W * (self.log_alpha < 8.0).float()) + self.bias if self.hasbias else F.linear(x, self.W * (self.log_alpha < 8.0).float()) + # return F.linear(x, self.W) + self.bias if self.hasbias else F.linear(x, self.W) + + def kl_reg(self): + # Return KL here -- a scalar + k1, k2, k3 = torch.Tensor([0.63576]).to(self.W.device), torch.Tensor([1.8732]).to(self.W.device), torch.Tensor([1.48695]).to(self.W.device) + kl = k1 * torch.sigmoid(k2 + k3 * self.log_alpha) - 0.5 * torch.log1p(torch.exp(-self.log_alpha)) + a = - torch.sum(kl) + return a \ No newline at end of file diff --git a/src/models/srgnn.py b/src/models/srgnn.py index 02c48e2..2f660eb 100644 --- a/src/models/srgnn.py +++ b/src/models/srgnn.py @@ -2,20 +2,26 @@ import torch as th import torch.nn as nn -import torch.nn.functional as F +import torch.nn.functional as F1 import dgl import dgl.ops as F import dgl.function as fn +from .sparsevd import LinearSVDO +from .gru import GRUCell + class SRGNNLayer(nn.Module): - def __init__(self, input_dim, output_dim, batch_norm=False, feat_drop=0.0, activation=None): + def __init__(self, input_dim, output_dim, batch_norm=False, feat_drop=0.0, threshold=3.0, activation=None, name=None): super().__init__() + self.name = name self.batch_norm = nn.BatchNorm1d(input_dim) if batch_norm else None self.dropout = nn.Dropout(feat_drop) - self.gru = nn.GRUCell(2 * input_dim, output_dim) - self.W1 = nn.Linear(input_dim, output_dim, bias=False) - self.W2 = nn.Linear(input_dim, output_dim, bias=False) + self.gru = GRUCell(2 * input_dim, output_dim, threshold=threshold, name=name+"_GRU") + # self.W1 = nn.Linear(input_dim, output_dim, bias=False) + # self.W2 = nn.Linear(input_dim, output_dim, bias=False) + self.W1 = LinearSVDO(input_dim, output_dim, threshold=threshold, bias=False, name=name+'_W1') + self.W2 = LinearSVDO(input_dim, output_dim, threshold=threshold, bias=False, name=name+'_W2') self.activation = activation def messager(self, edges): @@ -58,19 +64,19 @@ def __init__( output_dim, batch_norm=True, feat_drop=0.0, + threshold=3.0, activation=None, + name='AttnReadout' ): 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) + self.feat_drop = nn.Dropout(feat_drop) + self.fc_u = LinearSVDO(input_dim, hidden_dim, threshold=threshold, bias=False, name=name+'_fc_u') + self.fc_v = LinearSVDO(input_dim, hidden_dim, threshold=threshold, bias=True, name=name+'_fc_v') + self.fc_e = LinearSVDO(hidden_dim, 1, threshold=threshold, bias=False, name=name+'_fc_e') + self.fc_out = (nn.LinearSVDO(input_dim, output_dim, threshold=threshold, bias=False, name=name+'_fc_out') if output_dim != input_dim - else None - ) + else None) self.activation = activation def forward(self, g, feat, last_nodes): @@ -92,8 +98,10 @@ def forward(self, g, feat, last_nodes): class SRGNN(nn.Module): - def __init__(self, num_items, embedding_dim, num_layers, feat_drop=0.0): + def __init__(self, num_items, embedding_dim, num_layers, feat_drop=0.0, threshold=8, name="SRGNN"): super().__init__() + self.name = name + self.threshold = threshold self.embedding = nn.Embedding(num_items, embedding_dim) # self.indices = th.arange(num_items, dtype=th.long) self.register_buffer('indices', th.arange(num_items, dtype=th.long)) @@ -106,7 +114,9 @@ def __init__(self, num_items, embedding_dim, num_layers, feat_drop=0.0): input_dim, embedding_dim, batch_norm=None, - feat_drop=feat_drop + feat_drop=feat_drop, + threshold=threshold, + name=name + "_SRGNNLayer" + "_" + str(i) ) self.layers.append(layer) self.readout = AttnReadout( @@ -115,11 +125,13 @@ def __init__(self, num_items, embedding_dim, num_layers, feat_drop=0.0): embedding_dim, batch_norm=None, feat_drop=feat_drop, + threshold=threshold, activation=None, + name=name+"_AttnReadout" ) input_dim += embedding_dim self.feat_drop = nn.Dropout(feat_drop) - self.fc_sr = nn.Linear(input_dim, embedding_dim, bias=False) + self.fc_sr = LinearSVDO(input_dim, embedding_dim, threshold=threshold, bias=False, name=name+'_fc_sr') self.reset_parameters() @@ -131,18 +143,24 @@ def reset_parameters(self): def forward(self, mg, sg=None): iid = mg.ndata['iid'] feat = self.feat_drop(self.embedding(iid)) + # feat = F1.normalize(self.feat_drop(self.embedding(iid))) out = feat for i, layer in enumerate(self.layers): out = layer(mg, out) + feat = out + last_nodes = mg.filter_nodes(lambda nodes: nodes.data['last'] == 1) + # feat = F1.normalize(feat) 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) target = self.embedding(self.indices) + # sr = F1.normalize(sr) + # target = F1.normalize(target) logits = sr @ target.t() logits = th.log(nn.functional.softmax(logits, dim=-1)) return logits# , 0 diff --git a/src/scripts/main_msgifsr.py b/src/scripts/main_msgifsr.py index c48a4d2..0fbc0a7 100644 --- a/src/scripts/main_msgifsr.py +++ b/src/scripts/main_msgifsr.py @@ -148,12 +148,12 @@ def get_freer_gpu(): train_loader = DataLoader( train_set, batch_size=args.batch_size, - shuffle=True, + # shuffle=True, # drop_last=True, num_workers=args.num_workers, collate_fn=collate_fn, pin_memory=True, - # sampler=SequentialSampler(train_set) + sampler=SequentialSampler(train_set) ) test_loader = DataLoader( diff --git a/src/scripts/main_sparsevd_gnn.py b/src/scripts/main_sparsevd_gnn.py new file mode 100644 index 0000000..9e3e1d9 --- /dev/null +++ b/src/scripts/main_sparsevd_gnn.py @@ -0,0 +1,123 @@ +import argparse +import sys +import wandb + +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('--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=8, + 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), project='sparsevd') + +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_session_graph, + collate_fn_factory, +) +from src.utils.train import TrainRunner +from src.models import SRGNN + +dataset_dir = Path(args.dataset_dir) +print('reading dataset') +train_sessions, test_sessions, 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] + +train_set = AugmentedDataset(train_sessions) +test_set = AugmentedDataset(test_sessions) + + +collate_fn = collate_fn_factory(seq_to_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 = SRGNN(num_items, args.embedding_dim, args.num_layers, feat_drop=args.feat_drop) +device = th.device('cuda' 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') +mrr, hit = runner.train(args.epochs, args.log_interval) +print('MRR@20\tHR@20') +print(f'{mrr * 100:.3f}%\t{hit * 100:.3f}%') diff --git a/src/srgnn_sparsevd.pdf b/src/srgnn_sparsevd.pdf new file mode 100644 index 0000000..5e03ce6 Binary files /dev/null and b/src/srgnn_sparsevd.pdf differ diff --git a/src/utils/__pycache__/train.cpython-38.pyc b/src/utils/__pycache__/train.cpython-38.pyc index fd7111a..c36471b 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-37.pyc b/src/utils/data/__pycache__/collate.cpython-37.pyc deleted file mode 100644 index d7d6230..0000000 Binary files a/src/utils/data/__pycache__/collate.cpython-37.pyc and /dev/null differ diff --git a/src/utils/data/__pycache__/collate.cpython-38.pyc b/src/utils/data/__pycache__/collate.cpython-38.pyc index adce854..28a9a95 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 deleted file mode 100644 index a7596a4..0000000 Binary files a/src/utils/data/__pycache__/collate.cpython-39.pyc and /dev/null differ diff --git a/src/utils/data/__pycache__/dataset.cpython-37.pyc b/src/utils/data/__pycache__/dataset.cpython-37.pyc deleted file mode 100644 index 0acd391..0000000 Binary files a/src/utils/data/__pycache__/dataset.cpython-37.pyc and /dev/null differ diff --git a/src/utils/data/__pycache__/dataset.cpython-39.pyc b/src/utils/data/__pycache__/dataset.cpython-39.pyc deleted file mode 100644 index 73770ad..0000000 Binary files a/src/utils/data/__pycache__/dataset.cpython-39.pyc and /dev/null differ diff --git a/src/utils/data/__pycache__/preprocess.cpython-37.pyc b/src/utils/data/__pycache__/preprocess.cpython-37.pyc deleted file mode 100644 index 572d7bd..0000000 Binary files a/src/utils/data/__pycache__/preprocess.cpython-37.pyc and /dev/null differ diff --git a/src/utils/data/__pycache__/preprocess.cpython-38.pyc b/src/utils/data/__pycache__/preprocess.cpython-38.pyc deleted file mode 100644 index d791b54..0000000 Binary files a/src/utils/data/__pycache__/preprocess.cpython-38.pyc and /dev/null differ diff --git a/src/utils/train.py b/src/utils/train.py index 4de090a..cce6dc0 100644 --- a/src/utils/train.py +++ b/src/utils/train.py @@ -80,6 +80,7 @@ def __init__( self.epoch = 0 self.batch = 0 self.patience = patience + self.kl_weight = 0.02 def train(self, epochs, log_interval=100): max_mrr = 0 @@ -90,13 +91,24 @@ def train(self, epochs, log_interval=100): mrr, hit = evaluate(self.model, self.test_loader, self.device) for epoch in tqdm(range(epochs)): + self.model.train() - for batch in self.train_loader: + for idx, batch in enumerate(self.train_loader): + self.kl_weight = min(self.kl_weight+0.02, 1) inputs, labels = prepare_batch(batch, self.device) self.optimizer.zero_grad() scores = self.model(*inputs) assert not th.isnan(scores).any() loss = nn.functional.nll_loss(scores, labels) + + kl = 0.0 + for module in self.model.modules(): + if hasattr(module, 'kl_reg'): + kl = kl + module.kl_reg() + + loss *= len(self.train_loader) + loss += (kl * self.kl_weight) + loss.backward() self.optimizer.step() @@ -106,11 +118,18 @@ def train(self, epochs, log_interval=100): print(f'Batch {self.batch}: Loss = {mean_loss:.4f}, Time Elapsed = {time.time() - t:.2f}s') t = time.time() mean_loss = 0 + + # print('here') + # print(list(self.model.modules())) + for i, c in enumerate(self.model.modules()): + if hasattr(c, 'kl_reg'): + wandb.log({'sp_%s' % c.name: (c.log_alpha.data.cpu().numpy() < self.model.threshold).mean()}, step=self.batch) self.batch += 1 + self.scheduler.step() mrr, hit = evaluate(self.model, self.test_loader, self.device) - + # wandb.log({"hit": hit, "mrr": mrr}) print(f'Epoch {self.epoch}: MRR = {mrr * 100:.3f}%, Hit = {hit * 100:.3f}%')