explaining_framework/main.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#

import copy
import os
import time

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 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,
                                           obj_config_to_str, read_json,
                                           write_json)

# inference, time, force,


if __name__ == "__main__":
    args = parse_args()
    outline = ExplainingOutline(args.explaining_cfg_file)
    out_dir = os.path.join(
        outline.explaining_cfg.out_dir,
        outline.cfg.dataset.name,
        outline.model_signature,
    )
    makedirs(out_dir)

    dump_cfg(outline.cfg, os.path.join(out_dir, "config.yaml"))
    write_json(outline.model_info, os.path.join(out_dir, "info.json"))

    explainer_path = os.path.join(
        out_dir,
        outline.explaining_cfg.explainer.name,
        obj_config_to_str(outline.explaining_algorithm),
    )
    makedirs(explainer_path)
    dump_cfg(
        outline.explainer_cfg,
        os.path.join(explainer_path, "explainer_cfg.yaml"),
    )
    dump_cfg(
        outline.explaining_cfg,
        os.path.join(explainer_path, explaining_cfg.cfg_dest),
    )

    item, index = outline.get_item()
    while not (item is None or index is None):
        for attack in outline.attacks:
            attack_path = os.path.join(
                out_dir, attack.__class__.__name__, obj_config_to_str(attack)
            )
            makedirs(attack_path)
            data_attack_path = os.path.join(attack_path, f"{index}.json")
            data_attack = outline.get_attack(
                attack=attack, item=item, path=data_attack_path
            )
        item, index = outline.get_item()

    outline.reload_dataloader()
    makedirs(explainer_path)

    item, index = outline.get_item()
    while not (item is None or index is None):
        for attack in outline.attacks:
            attack_path_ = os.path.join(
                explainer_path, attack.__class__.__name__, obj_config_to_str(attack)
            )
            makedirs(attack_path_)
            data_attack_path_ = os.path.join(attack_path_, f"{index}.json")
            attack_data = outline.get_attack(
                attack=attack, item=item, path=data_attack_path_
            )
            exp = outline.get_explanation(item=attack_data, path=data_attack_path_)
            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
                    )
                    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
                        )
                        print("#################################")
                        print("Attack", attack.name)
                        print(
                            "ThresholdConfig",
                            "-".join([f"{k}={v}" for k, v in threshold_conf.items()]),
                        )
                        print("Metric", metric.name)
                        print("Val", out_metric)
                        print("Index", index)
                        print("#################################")

        item, index = outline.get_item()