Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| import neptune.new as neptune | |
| #from parser import parameter_parser | |
| from CatGCN.clustering import ClusteringMachine | |
| from CatGCN.clustergnn import ClusterGNNTrainer | |
| from CatGCN.utils import pos_preds_attr_distr, tab_printer, graph_reader, field_reader, target_reader, label_reader | |
| import time | |
| from CatGCN.fairness import Fairness | |
| def train_CatGCN(user_edge, user_field, user_gender, user_labels, seed, label, args): | |
| start_time = time.perf_counter() | |
| """ | |
| Parsing command line parameters, reading data, graph decomposition, fitting and scoring the model. | |
| """ | |
| #args = parameter_parser() | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed(seed) | |
| # tab_printer(args) | |
| graph = graph_reader(user_edge) | |
| field_index = field_reader(user_field) | |
| target = target_reader(user_gender) | |
| user_labels = label_reader(user_labels) | |
| print('args', args) | |
| # Instantiate Neptune client and log arguments | |
| print('token:', args.neptune_token) | |
| neptune_run = neptune.init( | |
| project = args.neptune_project, | |
| api_token = args.neptune_token, | |
| ) | |
| #neptune_run["sys/tags"].add(args.log_tags.split(",")) | |
| neptune_run["seed"] = seed | |
| #neptune_run["dataset"] = "JD-small" if "jd" in args.edge_path else "Alibaba-small" | |
| neptune_run["model"] = "CatGCN" | |
| neptune_run["label"] = label | |
| neptune_run["lr"] = args.lr | |
| neptune_run["L2"] = args.weight_decay | |
| neptune_run["dropout"] = args.dropout | |
| neptune_run["diag_probe"] = args.diag_probe | |
| neptune_run["nfm_units"] = args.nfm_units | |
| neptune_run["grn_units"] = args.grn_units | |
| neptune_run["gnn_hops"] = args.gnn_hops | |
| neptune_run["gnn_units"] = args.gnn_units | |
| neptune_run["balance_ratio"] = args.balance_ratio | |
| neptune_run["n_epochs"] = args.epochs | |
| clustering_machine = ClusteringMachine(args, graph, field_index, target) | |
| clustering_machine.decompose() | |
| gnn_trainer = ClusterGNNTrainer(args, clustering_machine, neptune_run) # todo add later neptune_run | |
| gnn_trainer.train_val_test() | |
| ## Compute accuracy per sensitive attribute group | |
| if args.dataset_name == 'nba' or args.dataset_name == 'pokec_z' or args.dataset_name == 'pokec_n': | |
| pos_preds_distr = pos_preds_attr_distr(user_labels, gnn_trainer.targets, gnn_trainer.predictions, clustering_machine.sg_test_nodes[0], label, args.sens_attr) | |
| else: | |
| pos_preds_distr = pos_preds_attr_distr(user_labels, gnn_trainer.targets, gnn_trainer.predictions, clustering_machine.sg_test_nodes[0], label, args.sens_attr) | |
| print(pos_preds_distr) | |
| #neptune_run["pos_preds_distr"] = pos_preds_distr | |
| ## Compute fairness metrics | |
| print("Fairness metrics on sensitive attributes '{}':".format(args.sens_attr)) | |
| fair_obj = Fairness(user_labels, clustering_machine.sg_test_nodes[0], gnn_trainer.targets, gnn_trainer.predictions, args.sens_attr, neptune_run, args.multiclass_pred, args.multiclass_sens) #todo add neptune_run later | |
| fair_obj.statistical_parity() | |
| fair_obj.equal_opportunity() | |
| fair_obj.overall_accuracy_equality() | |
| fair_obj.treatment_equality() | |
| elaps_time = (time.perf_counter() - start_time)/60 | |
| neptune_run["elaps_time"] = elaps_time | |
| neptune_run.stop() | |
| if __name__ == "__main__": | |
| train_CatGCN() | |