explaining_framework/main.py

128 lines
5.3 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
import copy
import logging
import os
import torch
from torch_geometric import seed_everything
from torch_geometric.data.makedirs import makedirs
from torch_geometric.explain import Explainer
from torch_geometric.explain.config import ThresholdConfig
from torch_geometric.graphgym.config import cfg
from torch_geometric.graphgym.utils.device import auto_select_device
from tqdm import tqdm
from explaining_framework.config.explaining_config import explaining_cfg
from explaining_framework.utils.explaining.cmd_args import parse_args
from explaining_framework.utils.explaining.outline import ExplainingOutline
from explaining_framework.utils.explanation.adjust import Adjust
from explaining_framework.utils.io import (dump_cfg, is_exists, notify_done,
obj_config_to_str, read_json,
set_printing, write_json)
if __name__ == "__main__":
args = parse_args()
outline = ExplainingOutline(args.explaining_cfg_file, args.gpu_id)
for attack in outline.attacks:
done_attack_raw = True
logging.info("Running %s: %s" % (attack.__class__.__name__, attack.name))
for item, index in tqdm(
zip(outline.dataset, outline.indexes), total=len(outline.dataset)
):
item = item.to(outline.cfg.accelerator)
attack_path = os.path.join(
outline.out_dir, attack.__class__.__name__, obj_config_to_str(attack)
)
makedirs(attack_path)
data_attack_path = os.path.join(attack_path, f"{index}.json")
if args.force_fastforward:
if os.path.exists(data_attack_path):
continue
data_attack = outline.get_attack(
attack=attack, item=item, path=data_attack_path
)
if data_attack is None:
done = False
continue
if done_attack_raw:
notify_done(attack_path)
for attack in outline.attacks:
logging.info("Running %s: %s" % (attack.__class__.__name__, attack.name))
logging.info(
"Running %s: %s"
% (outline.explainer.__class__.__name__, outline.explaining_algorithm.name),
)
done_attack_exp = True
done_exp = True
for item, index in tqdm(
zip(outline.dataset, outline.indexes), total=len(outline.dataset)
):
item = item.to(outline.cfg.accelerator)
attack_path_ = os.path.join(
outline.explainer_path,
attack.__class__.__name__,
obj_config_to_str(attack),
)
makedirs(attack_path_)
data_attack_path_ = os.path.join(attack_path_, f"{index}.json")
if args.force_fastforward:
if os.path.exists(data_attack_path_):
continue
attack_data = outline.get_attack(
attack=attack, item=item, path=data_attack_path_
)
if attack_data is None:
done_attack_exp = False
continue
exp = outline.get_explanation(item=attack_data, path=data_attack_path_)
if exp is None:
done_exp = False
continue
else:
for adjust in outline.adjusts:
adjust_path = os.path.join(
attack_path_,
adjust.__class__.__name__,
obj_config_to_str(adjust),
)
makedirs(adjust_path)
exp_adjust_path = os.path.join(adjust_path, f"{index}.json")
exp_adjust = outline.get_adjust(
adjust=adjust, item=exp, path=exp_adjust_path
)
for threshold_conf in outline.thresholds_configs:
outline.set_explainer_threshold_config(threshold_conf)
masking_path = os.path.join(
adjust_path,
"ThresholdConfig",
obj_config_to_str(threshold_conf),
)
makedirs(masking_path)
exp_masked_path = os.path.join(masking_path, f"{index}.json")
exp_masked = outline.get_threshold(
item=exp_adjust, path=exp_masked_path
)
if exp_masked is None:
done = False
continue
else:
for metric in outline.metrics:
metric_path = os.path.join(
masking_path,
metric.__class__.__name__,
obj_config_to_str(metric),
)
makedirs(metric_path)
metric_path = os.path.join(metric_path, f"{index}.json")
out_metric = outline.get_metric(
metric=metric, item=exp_masked, path=metric_path
)
if done_exp and done_exp:
notify_done(attack_path_)