def train(dataset, model, args, same_feat=True, val_dataset=None, test_dataset=None, writer=None, mask_nodes = True): writer_batch_idx = [0, 3, 6, 9] optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()), lr=0.001) iter = 0 best_val_result = { 'epoch': 0, 'loss': 0, 'acc': 0} test_result = { 'epoch': 0, 'loss': 0, 'acc': 0} train_accs = [] train_epochs = [] best_val_accs = [] best_val_epochs = [] test_accs = [] test_epochs = [] val_accs = [] for epoch in range(args.num_epochs): total_time = 0 avg_loss = 0.0 model.train() print('Epoch: ', epoch) for batch_idx, data in enumerate(dataset): begin_time = time.time() model.zero_grad() adj = Variable(data['adj'].float(), requires_grad=False).cuda() h0 = Variable(data['feats'].float(), requires_grad=False).cuda() label = Variable(data['label'].long()).cuda() batch_num_nodes = data['num_nodes'].int().numpy() if mask_nodes else None assign_input = Variable(data['assign_feats'].float(), requires_grad=False).cuda() ypred = model(h0, adj, batch_num_nodes, assign_x=assign_input) if not args.method == 'soft-assign' or not args.linkpred: loss = model.loss(ypred, label) else: loss = model.loss(ypred, label, adj, batch_num_nodes) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() iter += 1 avg_loss += loss #if iter % 20 == 0: # print('Iter: ', iter, ', loss: ', loss.data[0]) elapsed = time.time() - begin_time total_time += elapsed # log once per XX epochs if epoch % 10 == 0 and batch_idx == len(dataset) // 2 and args.method == 'soft-assign' and writer is not None: log_assignment(model.assign_tensor, writer, epoch, writer_batch_idx) if args.log_graph: log_graph(adj, batch_num_nodes, writer, epoch, writer_batch_idx, model.assign_tensor) avg_loss /= batch_idx + 1 if writer is not None: writer.add_scalar('loss/avg_loss', avg_loss, epoch) if args.linkpred: writer.add_scalar('loss/linkpred_loss', model.link_loss, epoch) print('Avg loss: ', avg_loss, '; epoch time: ', total_time) result = evaluate(dataset, model, args, name='Train', max_num_examples=100) train_accs.append(result['acc']) train_epochs.append(epoch) if val_dataset is not None: print("validation>>>>>>") val_result = evaluate(val_dataset, model, args, name='Validation') val_accs.append(val_result['acc']) if val_result['acc'] > best_val_result['acc'] - 1e-7: print("saving best results>>>>>") best_val_result['acc'] = val_result['acc'] best_val_result['epoch'] = epoch best_val_result['loss'] = avg_loss if test_dataset is not None: test_result = evaluate(test_dataset, model, args, name='Test') test_result['epoch'] = epoch if writer is not None: writer.add_scalar('acc/train_acc', result['acc'], epoch) writer.add_scalar('acc/val_acc', val_result['acc'], epoch) writer.add_scalar('loss/best_val_loss', best_val_result['loss'], epoch) if test_dataset is not None: writer.add_scalar('acc/test_acc', test_result['acc'], epoch) print('Best val result: ', best_val_result) best_val_epochs.append(best_val_result['epoch']) best_val_accs.append(best_val_result['acc']) if test_dataset is not None: print('Test result: ', test_result) test_epochs.append(test_result['epoch']) test_accs.append(test_result['acc']) matplotlib.style.use('seaborn') plt.switch_backend('agg') plt.figure() plt.plot(train_epochs, util.exp_moving_avg(train_accs, 0.85), '-', lw=1) if test_dataset is not None: plt.plot(best_val_epochs, best_val_accs, 'bo', test_epochs, test_accs, 'go') plt.legend(['train', 'val', 'test']) else: plt.plot(best_val_epochs, best_val_accs, 'bo') plt.legend(['train', 'val']) plt.savefig(gen_train_plt_name(args), dpi=600) plt.close() matplotlib.style.use('default') return model, val_accs
def train(dataset, model, user_feature, adjacency, args, same_feat=True, val_dataset=None, test_dataset=None, writer=None, mask_nodes=True): writer_batch_idx = [0, 3, 6, 9] optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001) iter = 0 best_val_result = {'epoch': 0, 'loss': 0, 'acc': 0} test_result = {'epoch': 0, 'loss': 0, 'acc': 0} train_accs = [] train_epochs = [] best_val_accs = [] best_val_epochs = [] test_accs = [] test_epochs = [] val_accs = [] adj = Variable(torch.Tensor(adjacency), requires_grad=False) h0 = Variable(torch.Tensor(user_feature), requires_grad=False) assign_input = h0 for epoch in range(args.num_epochs): begin_time = time.time() avg_loss = 0.0 model.train() print('Epoch: ', epoch) embed_all = [] assign_all = [] for batch_idx, data in enumerate(dataset): model.zero_grad() train_data = (data['uid'], data['iid']) batch_num_nodes = None # batch_num_nodes = data['num_nodes'].int().numpy() if mask_nodes else None # forward function invoking # logits = model(h0, adj, batch_num_nodes, assign_x=assign_input) # assign_x = h0 embed_all, assign_all = model(h0, adj, batch_num_nodes) if not args.method == 'soft-assign' or not args.linkpred: loss = model.pred_loss( model.pred_item_based_logits(embed_all, assign_all, train_data, args)) else: loss = model.pred_loss( model.pred_item_based_logits(embed_all, assign_all, train_data, args)) loss.backward() nn.utils.clip_grad_norm(model.parameters(), args.clip) optimizer.step() iter += 1 avg_loss += loss print('Avg loss: ', avg_loss / (batch_idx + 1)) ''' if epoch % 10 == 0 and batch_idx == len(dataset) // 2 and args.method == 'soft-assign' and writer is not None: log_assignment(model.assign_tensor, writer, epoch, writer_batch_idx) log_graph(adj, batch_num_nodes, writer, epoch, writer_batch_idx, model.assign_tensor) ''' # save item_embeds = model.get_item_embeds() save_numpy(embed_all, assign_all, item_embeds, epoch, args) avg_loss /= batch_idx + 1 elapsed = time.time() - begin_time if writer is not None: writer.add_scalar('loss/avg_loss', avg_loss, epoch) if args.linkpred: writer.add_scalar('loss/linkpred_loss', model.link_loss, epoch) print('Avg loss: ', avg_loss, '; epoch time: ', elapsed) if val_dataset is not None: val_result = evaluate(val_dataset, model, args, name='Validation') val_accs.append(val_result['acc']) if val_result['acc'] > best_val_result['acc'] - 1e-7: best_val_result['acc'] = val_result['acc'] best_val_result['epoch'] = epoch best_val_result['loss'] = avg_loss if test_dataset is not None: test_result = evaluate(test_dataset, model, args, name='Test') test_result['epoch'] = epoch print('Best val result: ', best_val_result) best_val_epochs.append(best_val_result['epoch']) best_val_accs.append(best_val_result['acc']) if test_dataset is not None: print('Test result: ', test_result) test_epochs.append(test_result['epoch']) test_accs.append(test_result['acc']) matplotlib.style.use('seaborn') plt.switch_backend('agg') plt.figure() plt.plot(train_epochs, util.exp_moving_avg(train_accs, 0.85), '-', lw=1) if test_dataset is not None: plt.plot(best_val_epochs, best_val_accs, 'bo', test_epochs, test_accs, 'go') plt.legend(['train', 'val', 'test']) else: plt.plot(best_val_epochs, best_val_accs, 'bo') plt.legend(['train', 'val']) plt.savefig(gen_train_plt_name(args), dpi=600) plt.close() matplotlib.style.use('default') return model, val_accs