def main(opt):
train_dataset = BADataset(opt.dataroot, opt.L, True, False, False)
train_dataloader = BADataloader(train_dataset, batch_size=opt.batchSize, \
shuffle=True, num_workers=opt.workers, drop_last=True)
valid_dataset = BADataset(opt.dataroot, opt.L, False, True, False)
valid_dataloader = BADataloader(valid_dataset, batch_size=opt.batchSize, \
shuffle=True, num_workers=opt.workers, drop_last=True)
test_dataset = BADataset(opt.dataroot, opt.L, False, False, True)
test_dataloader = BADataloader(test_dataset, batch_size=opt.batchSize, \
shuffle=True, num_workers=opt.workers, drop_last=True)
all_dataset = BADataset(opt.dataroot, opt.L, False, False, False)
all_dataloader = BADataloader(all_dataset, batch_size=opt.batchSize, \
shuffle=False, num_workers=opt.workers, drop_last=False)
opt.n_edge_types = train_dataset.n_edge_types
opt.n_node = train_dataset.n_node
net = STGCN(opt, kernel_size=2, n_blocks=1, state_dim_bottleneck=opt.state_dim, annotation_dim_bottleneck=opt.annotation_dim)
net.double()
print(net)
criterion = nn.MSELoss()
#criterion = nn.CosineSimilarity(dim=-1, eps=1e-6)
if opt.cuda:
net.cuda()
criterion.cuda()
optimizer = optim.Adam(net.parameters(), lr=opt.lr)
early_stopping = EarlyStopping(patience=opt.patience, verbose=True)
os.makedirs(OutputDir, exist_ok=True)
train_loss_ls = []
valid_loss_ls = []
test_loss_ls = []
for epoch in range(0, opt.niter):
train_loss = train(epoch, train_dataloader, net, criterion, optimizer, opt)
valid_loss = valid(valid_dataloader, net, criterion, opt)
test_loss = test(test_dataloader, net, criterion, opt)
train_loss_ls.append(train_loss)
valid_loss_ls.append(valid_loss)
test_loss_ls.append(test_loss)
early_stopping(valid_loss, net, OutputDir)
if early_stopping.early_stop:
print("Early stopping")
break
df = pd.DataFrame({'epoch':[i for i in range(1, len(train_loss_ls)+1)], 'train_loss': train_loss_ls, 'valid_loss': valid_loss_ls, 'test_loss': test_loss_ls})
df.to_csv(OutputDir + '/loss.csv', index=False)
net.load_state_dict(torch.load(OutputDir + '/checkpoint.pt'))
inference(all_dataloader, net, criterion, opt, OutputDir)
{
'num_of_time_conv_filters1': 32,
'num_of_time_conv_filters2': 64,
'K_t': 3,
'num_of_cheb_filters': 32,
'K': 1,
'cheb_polys': cheb_polys
},
{
'num_of_time_conv_filters1': 32,
'num_of_time_conv_filters2': 128,
'K_t': 3,
'num_of_cheb_filters': 32,
'K': 1,
'cheb_polys': cheb_polys
}
]
net = STGCN(backbones, 128)
net.initialize(ctx = ctx)
loss_function = gluon.loss.L2Loss()
trainer = Trainer(net.collect_params(), optimizer, {'learning_rate': learning_rate})
training_dataloader = gluon.data.DataLoader(gluon.data.ArrayDataset(training_data_norm, training_target), batch_size = batch_size, shuffle = True)
validation_dataloader = gluon.data.DataLoader(gluon.data.ArrayDataset(val_data_norm, val_target), batch_size = batch_size, shuffle = False)
testing_dataloader = gluon.data.DataLoader(gluon.data.ArrayDataset(testing_data_norm, testing_target), batch_size = batch_size, shuffle = False)
if not os.path.exists('stgcn_params'):
os.mkdir('stgcn_params')
train_loss_list, val_loss_list, test_loss_list = train_model(net, training_dataloader, validation_dataloader, testing_dataloader)
def main(dataset):
batch_size = 1024
graph = Graph(dataset)
if dataset == 'meituan':
data_size = 6
has_user_region=True
else:
data_size = 5
has_user_region = False
g = graph.g
g.readonly()
embeddings = graph.embeddings
num_nodes = graph.g.number_of_nodes()
model = Recommender(STGCN(num_nodes, 64, 27, 2, None, embeddings))
model.cuda()
opt = torch.optim.Adam(model.parameters(), lr=0.0001, weight_decay=0)
model.train()
print('train neighbors')
for i in range(30):
total_loss = 0
g.apply_edges(edge_func1)
neighbors = graph.neighbors
neg_neighbors = np.random.choice(range(num_nodes), neighbors.shape[0], replace=True).reshape(-1, 1)
neighbors_data = np.concatenate((neighbors, neg_neighbors), axis=1)
data = torch.from_numpy(neighbors_data).cuda()
seed_nodes = data.reshape(-1)
batches = data.split(batch_size)
sampler = dgl.contrib.sampling.NeighborSampler(
g,
batch_size * 3,
5,
2,
seed_nodes=seed_nodes,
num_workers=11)
count = 0
for batch, nf in zip(batches, sampler):
nf.copy_from_parent(ctx=device)
batch_nid = nf.map_from_parent_nid(-1, batch.reshape(-1), True)
batch_nid = batch_nid.reshape(-1, 3).cuda()
loss = model.train_region(nf, batch_nid)
opt.zero_grad()
loss.backward()
total_loss += loss.item()
opt.step()
count += 1
print('loss', total_loss / count)
for epoch in range(300):
model.train()
begin = time.time()
total_loss = 0
count = 0
for t in range(0, 24, 2):
g.apply_edges(partial(edge_func, t))
pos = graph.time_train[t]
neg_pois = graph.pid2rid[['p', 'r']].sample(n=pos.shape[0], replace=True).to_numpy(copy=True)
data = np.concatenate((pos, neg_pois), axis=1)
data.astype(np.int)
data = torch.from_numpy(data).cuda()
seed_nodes = data.reshape(-1)
batches = data.split(batch_size)
sampler = dgl.contrib.sampling.NeighborSampler(
g,
batch_size * data_size,
5,
2,
seed_nodes=seed_nodes,
transition_prob='final_weight',
prefetch=False,
num_workers=11)
for batch, nf in zip(batches, sampler):
nf.copy_from_parent(ctx=device)
batch_nid = nf.map_from_parent_nid(-1, batch.reshape(-1), True)
batch_nid = batch_nid.reshape(-1, data_size).cuda()
loss = model(nf, batch_nid, has_user_region=has_user_region)
opt.zero_grad()
loss.backward()
total_loss += loss.item()
opt.step()
count += 1
print('epoch:{}, loss:{}, time:{}'.format(epoch, total_loss / count, time.time() - begin))
if epoch % 20 ==0 and epoch != 0:
model.eval()
with torch.no_grad():
recallk(graph, model, has_user_region=has_user_region)
model.eval()
with torch.no_grad():
recallk(graph, model, has_user_region=has_user_region)
