DP-GNN design that ensures both model weights and inference procedure differentially private (NeurIPS 2023)
import logging
import warnings
from rich.traceback import install
from rich.logging import RichHandler
from core.console.console import Console

globals = {
    'debug': False,
}

# define main and error consoles
console = Console(tab_size=4)
error_console = Console(stderr=True, tab_size=4)

# setup console for tracebacks
install(console=error_console, width=error_console.width)

# create logger
log_handler = RichHandler(
    console=console, 
    omit_repeated_times=True,
    log_time_format="[%X]"
)
logger = logging.getLogger('gap')
logger.setLevel(logging.INFO)
logger.addHandler(log_handler)

# setup warnings
logging.getLogger("py.warnings").addHandler(log_handler)
logging.getLogger("py.warnings").propagate = False
logging.captureWarnings(True)
warnings.simplefilter("ignore")

core/__init__.py
thupchnsky/dp-gnn
from typing import Annotated, Callable, Literal, Union, get_args, get_origin
from core import console
import inspect
from argparse import SUPPRESS, ArgumentParser, ArgumentTypeError, Namespace
from core.utils import dict2table


ArgType = Union[Namespace, dict[str, object]]
ArgInfo = dict[str, object]


class ArgWithLiteral:
    def __init__(self, main_type, literals):
        self.main_type = main_type
        self.literals = literals

    def __call__(self, arg):
        try:
            return self.main_type(arg)
        except ValueError:
            if arg in self.literals:
                return arg
            else:
                raise ArgumentTypeError(f'{arg} is not a valid literal')

def boolean(v: Union[str, bool]) -> bool:
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1', 'on'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0', 'off'):
        return False
    else:
        raise ArgumentTypeError('Boolean value expected.')


def remove_prefix(kwargs: dict, prefix: str) -> dict:
    return {k[len(prefix):]: v for k, v in kwargs.items() if k.startswith(prefix)}


def strip_kwargs(callable: Callable, kwargs: dict, prefix: str='') -> dict[str, object]:
    signature = inspect.signature(callable)
    parameters = signature.parameters
    out_kwargs = {}
    for arg, value in kwargs.items():
        if arg.startswith(prefix) and arg[len(prefix):] in parameters:
            out_kwargs[arg] = value

    # check if the function has kwargs
    for _, param in parameters.items():
        annotation = param.annotation
        if get_origin(annotation) is Annotated:
            metadata = get_args(annotation)[1]
            bases = metadata.get('bases', [])
            prefixes = metadata.get('prefixes', [''] * len(bases))
            for base_callable, pr in zip(bases, prefixes):
                out_kwargs.update(strip_kwargs(base_callable, kwargs, prefix=prefix+pr))
        elif param.kind == inspect.Parameter.VAR_KEYWORD:
            return kwargs

    return out_kwargs


def create_arguments(callable: Callable, parser: ArgumentParser, exclude: list = [], prefix: str = ''):
    arguments_added = [action.dest for action in parser._actions]
    parameters = inspect.signature(callable).parameters

    # iterate over the parameters
    for param_name, param_obj in parameters.items():
        arg_name = prefix + param_name

        # skip the parameters that are in the exclude list
        if param_name in exclude or arg_name in arguments_added:
            continue
        
        # get the annotation
        annot_obj = param_obj.annotation

        # only process annotated parameters
        if get_origin(annot_obj) is Annotated:

            # extract parameter type and metadata from annotation
            annotation = get_args(annot_obj)
            metadata: dict = annotation[1]
            param_type = metadata.get('type', annotation[0])

            # get the base callable arguments
            bases = metadata.get('bases', False)

            if bases:
                # if there are base callables, recursively add their args to the parser
                prefixes = metadata.get('prefixes', [''] * len(bases))
                for base_callable, pr in zip(bases, prefixes):
                    create_arguments(
                        callable=base_callable, 
                        parser=parser, 
                        exclude=metadata.get('exclude', []) + exclude,
                        prefix=prefix+pr
                    )
            else:
                # if there are no base callables, add the parameter to the parser
                metadata['type'] = param_type
                metadata['dest'] = arg_name

                # if the parameter has a default value, add it to the parser
                # otherwise, set the parameter as required
                if param_obj.default is not inspect.Parameter.empty:
                    metadata['default'] = param_obj.default
                else:
                    metadata['required'] = True
                    metadata['default'] = SUPPRESS
                    metadata['help'] += ' (required)'

                # tweak specific data types
                if param_type is bool:
                    # process boolean parameters
                    metadata['type'] = boolean
                    metadata['nargs'] = '?'
                    metadata['const'] = True
                elif get_origin(param_type) is Union:
                    # process union parameters
                    sub_types = get_args(param_type)
                    if len(sub_types) == 2 and get_origin(sub_types[0]) is Literal:
                        metadata['type'] = ArgWithLiteral(main_type=sub_types[1], literals=get_args(sub_types[0]))
                        metadata['metavar'] = f"<{sub_types[1].__name__}>" + '|{' +  ','.join(map(str, get_args(sub_types[0]))) + '}'

                # if metadata contains "choices", the parser uses that as meta variable
                # otherwise, if "metavar" is not provided, set the meta variable to  <parameter type>
                if 'choices' not in metadata:
                    try:
                        metadata['metavar'] = metadata.get('metavar', f'<{param_type.__name__}>')
                    except: pass

                # create options based on parameter name
                options = {f'--{arg_name}', f'--{arg_name.replace("_", "-")}'}

                # add custome options if provided
                custom_options = metadata.pop('option', [])
                custom_options = [custom_options] if isinstance(custom_options, str) else custom_options
                for option in custom_options:
                    idx = max([i for i,c in enumerate(option) if c == '-'])
                    option = option[:idx+1] + prefix + option[idx+1:]
                    options.add(option)

                # sort option names based on their length
                options = sorted(sorted(list(options)), key=len)

                # add the parameter to the parser
                parser.add_argument(*options, **metadata)


def print_args(args: ArgType, num_cols: int = 4):
    args = args if isinstance(args, dict) else vars(args)
    table = dict2table(args, num_cols=num_cols, title='program arguments')
    console.print()
    console.info(table)
    console.print()

core/args/utils.py
thupchnsky/dp-gnn
import os
import logging
from rich.console import Console as RichConsole
from rich.logging import RichHandler
from typing import Optional, Union
from rich.console import JustifyMethod, NewLine
from rich.scope import render_scope
from rich.segment import Segment
from rich.styled import Styled
from rich.style import Style
from core.console.status import LogStatus

class Console(RichConsole):
    CRITICAL = logging.CRITICAL
    FATAL = logging.FATAL
    ERROR = logging.ERROR
    WARNING = logging.WARNING
    WARN = logging.WARN
    INFO = logging.INFO
    DEBUG = logging.DEBUG
    NOTSET = logging.NOTSET

    def __init__(self, *args, log_level: int = INFO, **kwrags):
        super().__init__(*args, **kwrags)
        self.log_level = log_level

    def status (self, status: str, level: int=INFO) -> LogStatus:
        return LogStatus(status, console=self, level=level, enabled=level >= self.log_level)

    def log(
        self,
        *objects,
        level: int = logging.INFO,
        sep: str = " ",
        end: str = "\n",
        style: Union[str, Style] = None,
        justify: Optional[JustifyMethod] = None,
        emoji: Optional[bool] = None,
        markup: Optional[bool] = None,
        highlight: Optional[bool] = None,
        log_locals: bool = False,
        _stack_offset: int = 2,
    ) -> None:
        """Log rich content to the terminal.
        Args:
            objects (positional args): Objects to log to the terminal.
            sep (str, optional): String to write between print data. Defaults to " ".
            end (str, optional): String to write at end of print data. Defaults to "\\\\n".
            style (Union[str, Style], optional): A style to apply to output. Defaults to None.
            justify (str, optional): One of "left", "right", "center", or "full". Defaults to ``None``.
            emoji (Optional[bool], optional): Enable emoji code, or ``None`` to use console default. Defaults to None.
            markup (Optional[bool], optional): Enable markup, or ``None`` to use console default. Defaults to None.
            highlight (Optional[bool], optional): Enable automatic highlighting, or ``None`` to use console default. Defaults to None.
            log_locals (bool, optional): Boolean to enable logging of locals where ``log()``
                was called. Defaults to False.
            _stack_offset (int, optional): Offset of caller from end of call stack. Defaults to 1.
        """
        if level < self.log_level:
            return

        if not objects:
            objects = (NewLine(),)

        render_hooks = self._render_hooks[:]

        with self:
            renderables = self._collect_renderables(
                objects, sep, end, justify=justify, emoji=emoji, markup=markup, highlight=highlight
            )
            if style is not None:
                renderables = [Styled(renderable, style) for renderable in renderables]

            filename, line_no, locals = self._caller_frame_info(_stack_offset)
            link_path = None if filename.startswith("<") else os.path.abspath(filename)
            path = filename.rpartition(os.sep)[-1]
            if log_locals:
                locals_map = {key: value for key, value in locals.items() if not key.startswith("__")}
                renderables.append(render_scope(locals_map, title="[i]locals"))

            record = logging.LogRecord(name=None, level=level, pathname=None, lineno=None, msg=None, args=None, exc_info=None)
            handler = RichHandler(console=self)
            self._log_render.show_level = True

            renderables = [
                self._log_render(
                    self,
                    renderables,
                    level=handler.get_level_text(record),
                    log_time=self.get_datetime(),
                    path=path,
                    line_no=line_no,
                    link_path=link_path,
                )
            ]
            for hook in render_hooks:
                renderables = hook.process_renderables(renderables)
            new_segments = []
            extend = new_segments.extend
            render = self.render
            render_options = self.options
            for renderable in renderables:
                extend(render(renderable, render_options))
            buffer_extend = self._buffer.extend
            for line in Segment.split_and_crop_lines(new_segments, self.width, pad=False):
                buffer_extend(line)

    def debug(self, *args, **kwargs):
        self.log(*args, level=self.DEBUG, **kwargs)

    def info(self, *args, **kwargs):
        self.log(*args, level=self.INFO, **kwargs)

    def warning(self, *args, **kwargs):
        self.log(*args, level=self.WARNING, **kwargs)

    def error(self, *args, **kwargs):
        self.log(*args, level=self.ERROR, **kwargs)

core/console/console.py
thupchnsky/dp-gnn
from rich.console import Console as RichConsole
from rich.logging import RichHandler
from rich.spinner import Spinner
from rich.table import Table
from rich.status import Status
from rich.live import Live
from rich._log_render import LogRender
from time import time
import logging


class LogStatus(Status):
    def __init__(self,
        status,
        console: RichConsole,
        level: int = logging.INFO,
        enabled: bool = True,
        speed: float = 1.0,
        refresh_per_second: float = 12.5,
    ):
        super().__init__(status, 
            console=console, 
            spinner='simpleDots', 
            speed=speed, 
            refresh_per_second=refresh_per_second
        )
        
        self.status = status
        self.level = level
        self.enabled = enabled
        spinner = Spinner('simpleDots', style='status.spinner', speed=speed)
        record = logging.LogRecord(name=None, level=level, pathname=None, lineno=None, msg=None, args=None, exc_info=None)
        handler = RichHandler(console=console)
        table = Table.grid()
        table.add_row(self.status, spinner)
        
        self._spinner = LogRender(show_level=True, time_format='[%X]')(
            console=console, 
            level=handler.get_level_text(record),
            renderables=[table]
        )
        self._live = Live(
            self.renderable,
            console=console,
            refresh_per_second=refresh_per_second,
            transient=True,
        )
        
    def __enter__(self):
        if self.enabled:
            self._start_time = time()
            return super().__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.enabled:
            super().__exit__(exc_type, exc_val, exc_tb)
            self._end_time = time()
            self.console.log(f'{self.status}...done in {self._end_time - self._start_time:.2f} s', level=self.level)

core/console/status.py
thupchnsky/dp-gnn
from core.data.loader.node import NodeDataLoader
from core.data.loader.adj import AdjDataLoader
from core.data.loader.adj import SparseAdjRowLoader
core/data/loader/__init__.py
thupchnsky/dp-gnn
from typing import Literal, Optional, Union
import torch
import numpy as np
from collections.abc import Iterator
from torch.utils.data import Dataset, DataLoader
from torch_geometric.data import Data
from torch_geometric.utils import to_undirected, sort_edge_index
from torch_geometric.transforms import ToSparseTensor
from core.modules.base import Stage
from torch_sparse import SparseTensor


class SparseDataset(Dataset):

    def __init__(self, 
                 data: Data,
                 stage: Stage):
        
        self.device = data.x.device
        self.adj_mat = data.adj_t.t().to_dense()[data[f'{stage}_mask']]
        self.label = data.y[data[f'{stage}_mask']]

    def __len__(self):
        return self.adj_mat.shape[0]

    def __getitem__(self, idx):
        return self.adj_mat[idx], self.label[idx]


def AdjDataLoader(data: Data, 
                  stage: Stage,
                  batch_size: Union[int, Literal['full']] = 'full', 
                  shuffle: bool = True, 
                  drop_last: bool = False, 
                  sampling: str = 'None'):
    
    assert sampling == 'None', 'Other sampling method is not supported for this loader!'
    
    if batch_size == 'full':
        node_indices = data[f'{stage}_mask'].nonzero().view(-1)
        batch_size = node_indices.size(0)
    
    dataset = SparseDataset(data, stage)
    
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)


class SparseAdjRowLoader:
    """
    Args:
        data (Data): The graph data object.
        stage (Stage): Training stage. One of 'train', 'val', 'test'.
        batch_size (int or 'full', optional): The batch size.
            If set to 'full', the entire graph is used as a single batch.
            (default: 'full')
        hops (int, optional): The number of hops to sample neighbors.
            If set to None, all neighbors are included. (default: None)
        shuffle (bool, optional): If set to True, the nodes are shuffled
            before batching. (default: True)
        drop_last (bool, optional): If set to True, the last batch is
            dropped if it is smaller than the batch size. (default: False)
        poisson_sampling (bool, optional): If set to True, poisson sampling
            is used to sample nodes. (default: False)
        wor_sampling (bool, optional): If set to True, a random yet fixed size
            sampling is used to sample nodes. (default: False)
    """
    def __init__(self, 
                 data: Data, 
                 stage: Stage,
                 batch_size: Union[int, Literal['full']] = 'full',
                 shuffle: bool = True, 
                 drop_last: bool = False, 
                 poisson_sampling: bool = False,
                 wor_sampling: bool = False):

        self.adj = data.adj_t.t()
        self.label = data.y
        self.stage = stage
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.poisson_sampling = poisson_sampling
        self.wor_sampling = wor_sampling
        self.device = data.x.device
        self.node_indices = data[f'{stage}_mask'].nonzero().view(-1)
        self.num_nodes = self.node_indices.size(0)
        
        assert not (self.poisson_sampling and self.wor_sampling)
        
    def __iter__(self) -> Iterator[Data]:
        if self.batch_size == 'full':
            row, col, _ = self.adj[self.node_indices, :].coo()
            y = self.label[self.node_indices]
            
            # processing the row
            row[1:] = torch.diff(row)
            row[0] = 0
            row = row.to(torch.bool)
            row = torch.cumsum(row,0)
            
            yield SparseTensor(row=row, col=col, sparse_sizes=(self.num_nodes, self.label.size(0))).to_torch_sparse_coo_tensor().to(self.device), y.to(self.device)
            return

        if self.shuffle and not (self.poisson_sampling or self.wor_sampling):
            perm = torch.randperm(self.num_nodes, device=self.device)
            self.node_indices = self.node_indices[perm]

        for i in range(0, self.num_nodes, self.batch_size):
            if self.drop_last and i + self.batch_size > self.num_nodes:
                break

            if self.poisson_sampling:
                sampling_prob = self.batch_size / self.num_nodes
                sample_mask = torch.rand(self.num_nodes, device=self.device) < sampling_prob
                batch_nodes = self.node_indices[sample_mask]
            elif self.wor_sampling:
                batch_nodes = self.node_indices[torch.randperm(self.num_nodes)[:self.batch_size]]
            else:    
                batch_nodes = self.node_indices[i: min(i + self.batch_size, self.num_nodes)]
            
            
            row, col, _ = self.adj[batch_nodes, :].coo()
            y = self.label[batch_nodes]
            
            # processing the row
            row[1:] = torch.diff(row)
            row[0] = 0
            row = row.to(torch.bool)
            row = torch.cumsum(row,0)
            
            yield SparseTensor(row=row, col=col, sparse_sizes=(batch_nodes.size(0), self.label.size(0))).to_dense().to(self.device), y.to(self.device)

    def __len__(self) -> int:
        if self.batch_size == 'full':
            return 1
        elif self.drop_last:
            return self.num_nodes // self.batch_size
        else:
            return (self.num_nodes + self.batch_size - 1) // self.batch_size

core/data/loader/adj.py
thupchnsky/dp-gnn
from typing import Literal, Optional, Union
import torch
from collections.abc import Iterator
from torch_geometric.data import Data
from torch_geometric.utils import k_hop_subgraph
from torch_geometric.transforms import ToSparseTensor
from core.modules.base import Stage


class NodeDataLoader:
    """ A fast dataloader for node-wise training.

    Args:
        data (Data): The graph data object.
        stage (Stage): Training stage. One of 'train', 'val', 'test'.
        batch_size (int or 'full', optional): The batch size.
            If set to 'full', the entire graph is used as a single batch.
            (default: 'full')
        hops (int, optional): The number of hops to sample neighbors.
            If set to None, all neighbors are included. (default: None)
        shuffle (bool, optional): If set to True, the nodes are shuffled
            before batching. (default: True)
        drop_last (bool, optional): If set to True, the last batch is
            dropped if it is smaller than the batch size. (default: False)
        poisson_sampling (bool, optional): If set to True, poisson sampling
            is used to sample nodes. (default: False)
        wor_sampling (bool, optional): If set to True, a random yet fixed size
            sampling is used to sample nodes. (default: False)
    """
    def __init__(self, 
                 data: Data, 
                 stage: Stage,
                 batch_size: Union[int, Literal['full']] = 'full', 
                 hops: Optional[int] = None,
                 shuffle: bool = True, 
                 drop_last: bool = False, 
                 poisson_sampling: bool = False,
                 wor_sampling: bool = False):

        self.data = data
        self.stage = stage
        self.batch_size = batch_size
        self.hops = hops
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.poisson_sampling = poisson_sampling
        self.wor_sampling = wor_sampling
        self.device = data.x.device
        
        assert not (self.poisson_sampling and self.wor_sampling)

        if batch_size != 'full':
            self.node_indices = data[f'{stage}_mask'].nonzero().view(-1)
            self.num_nodes = self.node_indices.size(0)

    def __iter__(self) -> Iterator[Data]:
        if self.batch_size == 'full':
            yield self.data
            return

        if self.shuffle and not (self.poisson_sampling or self.wor_sampling):
            perm = torch.randperm(self.num_nodes, device=self.device)
            self.node_indices = self.node_indices[perm]

        for i in range(0, self.num_nodes, self.batch_size):
            if self.drop_last and i + self.batch_size > self.num_nodes:
                break

            if self.poisson_sampling:
                sampling_prob = self.batch_size / self.num_nodes
                sample_mask = torch.rand(self.num_nodes, device=self.device) < sampling_prob
                batch_nodes = self.node_indices[sample_mask]
            elif self.wor_sampling:
                batch_nodes = self.node_indices[torch.randperm(self.num_nodes)[:self.batch_size]]
            else:    
                batch_nodes = self.node_indices[i: min(i + self.batch_size, self.num_nodes)]

            if self.hops is None:
                batch_mask = torch.zeros(self.data.num_nodes, device=self.device, dtype=torch.bool)
                batch_mask[batch_nodes] = True

                data = Data(**self.data.to_dict())
                data[f'{self.stage}_mask'] = data[f'{self.stage}_mask'] & batch_mask
            else:
                if not hasattr(self, 'edge_index'):
                    self.edge_index = torch.stack(self.data.adj_t.t().coo()[:2], dim=0)

                subset, batch_edge_index, mapping, _ = k_hop_subgraph(
                    node_idx=batch_nodes, 
                    num_hops=self.hops, 
                    edge_index=self.edge_index, 
                    relabel_nodes=True, 
                    num_nodes=self.data.num_nodes
                )

                batch_mask = torch.zeros(subset.size(0), device=self.device, dtype=torch.bool)
                batch_mask[mapping] = True

                data = Data(
                    x=self.data.x[subset],
                    y=self.data.y[subset],
                    edge_index=batch_edge_index,
                )
                data[f'{self.stage}_mask'] = batch_mask
                data = ToSparseTensor()(data)
            
            yield data
            
    def __len__(self) -> int:
        if self.batch_size == 'full':
            return 1
        elif self.drop_last:
            return self.num_nodes // self.batch_size
        else:
            return (self.num_nodes + self.batch_size - 1) // self.batch_size

core/data/loader/node.py
thupchnsky/dp-gnn
from core.data.transforms.filter_class import FilterClassByCount
from core.data.transforms.bound_degree import BoundDegree, BoundOutDegree
from core.data.transforms.remove_self_loops import RemoveSelfLoops
from core.data.transforms.split import RandomDataSplit
from core.data.transforms.remove_isolated_nodes import RemoveIsolatedNodes
from core.data.transforms.add_self_loops import AddSelfLoops
from core.data.transforms.normalize_degree import NormalizeDegree
core/data/transforms/__init__.py
thupchnsky/dp-gnn
from torch_geometric.utils import add_self_loops
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data
from torch_sparse import eye, SparseTensor


class AddSelfLoops(BaseTransform):
    def __init__(self, device: str):
        self.device = device
        
    def __call__(self, data: Data) -> Data:
        if hasattr(data, 'edge_index') and data.edge_index is not None:
            data.edge_index, _ = add_self_loops(data.edge_index)
        elif hasattr(data, 'adj_t'):
            N = data.adj_t.size(0)
            temp, _ = eye(N)
            row, col = temp
            I = SparseTensor(row=row, col=col).to(self.device)
            data.adj_t = data.adj_t + I
        else:
            raise TypeError("Not supported data type for edges!")
        return data

core/data/transforms/add_self_loops.py
thupchnsky/dp-gnn
import torch
import torch.utils.cpp_extension
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data
from torch_sparse import SparseTensor
from torch_geometric.loader.utils import filter_data
from torch_geometric.sampler.utils import to_csc

from torch_geometric.utils import is_undirected, to_undirected, to_scipy_sparse_matrix, from_scipy_sparse_matrix
import scipy.sparse as sp


class BoundOutDegree(BaseTransform):
    def __init__(self, max_out_degree: int):
        self.num_neighbors = max_out_degree
        self.with_replacement = False

    def __call__(self, data: Data) -> Data:
        data.adj_t = data.adj_t.t()
        data = self.sample(data)
        data.adj_t = data.adj_t.t()
        return data

    def sample(self, data: Data) -> Data:
        colptr, row, perm = to_csc(data, device='cpu')
        index = torch.arange(0, data.num_nodes-1, dtype=int)
        sample_fn = torch.ops.torch_sparse.neighbor_sample
        node, row, col, edge = sample_fn(
            colptr, row, index, [self.num_neighbors], self.with_replacement, True
        )
        data = filter_data(data, node, row, col, edge, perm)
        return data


class BoundDegree(BaseTransform):
    """
    We need a slight modification to make it work correctly.
    """
    
    def __init__(self, max_degree: int):
        self.max_deg = max_degree
        torch.ops.load_library("csrc/sampler.so")
        try:
            edge_sampler = torch.ops.my_ops.sample_edge
        except (AttributeError, RuntimeError):
            torch.utils.cpp_extension.load(
                name="sampler",
                sources=['csrc/sampler.cpp'],
                build_directory='csrc',
                is_python_module=False,
                verbose=False,
            )
            edge_sampler = torch.ops.my_ops.sample_edge
        
        self.edge_sampler = edge_sampler

    def __call__(self, data: Data) -> Data:
        N = data.num_nodes
        adj: SparseTensor = data.adj_t.t()
        device = adj.device()
        
        # Here, we need to first ensure that adj is upper triangle matrix.
        # Otherwise degree computation is incorrect and resulting adj can be assymmetric
        row, col, _ = adj.coo()
        edge_index = torch.stack([row,col],dim=0)
        
        # Ensure it is undirected
        if not is_undirected(edge_index):
            edge_index = to_undirected(edge_index)
        adj = to_scipy_sparse_matrix(edge_index)
        
        # get only the upper triangle part
        adj = sp.triu(adj,1)
        edge_index, _ = from_scipy_sparse_matrix(adj)
        row, col = edge_index
        E = len(row)
        
        # Below are the same, except at the end we need to make adj undirected again
        perm = torch.randperm(E)
        row, col = row[perm], col[perm]
        row, col = self.edge_sampler(row.tolist(), col.tolist(), N, self.max_deg)
        adj = SparseTensor(row=row, col=col).to(device)
        
        # make adj undirected
        data.adj_t = adj.t() + adj
        return data

core/data/transforms/bound_degree.py
thupchnsky/dp-gnn
import torch
import torch.nn.functional as F
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data


class FilterClassByCount(BaseTransform):
    def __init__(self, min_count: int, remove_unlabeled=False):
        self.min_count = min_count
        self.remove_unlabeled = remove_unlabeled

    def __call__(self, data: Data) -> Data:
        assert hasattr(data, 'y')

        y: torch.Tensor = F.one_hot(data.y)
        counts = y.sum(dim=0)
        y = y[:, counts >= self.min_count]
        mask = y.sum(dim=1).bool()        # nodes to keep
        data.y = y.argmax(dim=1)

        if self.remove_unlabeled:
            data = data.subgraph(mask)
        else:
            data.y[~mask] = -1                # set filtered nodes as unlabeled
            if hasattr(data, 'train_mask'):
                data.train_mask = data.train_mask & mask
                data.val_mask = data.val_mask & mask
                data.test_mask = data.test_mask & mask

        return data
core/data/transforms/filter_class.py
thupchnsky/dp-gnn
from torch_geometric.utils import add_self_loops
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data
from torch_sparse import SparseTensor


class NormalizeDegree(BaseTransform):
    """
    Perform D^{-beta}AD^{beta-1} degree normalization 
    """
    def __init__(self, beta=1.0):
        self.beta = beta
        
    def __call__(self, data: Data) -> Data:
        if hasattr(data, 'edge_index') and data.edge_index is not None:
            # Below implementation is for similicity. In practice, you should always use SparseTensor
            row, col = data.edge_index
            adj = SparseTensor(row=row,col=col)
            deg = adj.sum(dim=1)
            deg_inv_L = deg.pow(-self.beta)
            deg_inv_R = deg.pow(-(1-self.beta))
            deg_inv_L[deg_inv_L == float('inf')] = 0
            deg_inv_R[deg_inv_R == float('inf')] = 0
            adj = deg_inv_L.view(-1, 1) * adj * deg_inv_R.view(1, -1)
            row, col, _ = adj.coo()
            data.edge_index = torch.stack([row,col],dim=0)
            
        if hasattr(data, 'adj_t'):
            adj = data.adj_t.t()
            deg = adj.sum(dim=1)
            deg_inv_L = deg.pow(-self.beta)
            deg_inv_R = deg.pow(-(1-self.beta))
            deg_inv_L[deg_inv_L == float('inf')] = 0
            deg_inv_R[deg_inv_R == float('inf')] = 0
            adj = deg_inv_L.view(-1, 1) * adj * deg_inv_R.view(1, -1)
            data.adj_t = adj.t()
            
        return data

core/data/transforms/normalize_degree.py
thupchnsky/dp-gnn
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data


class RemoveIsolatedNodes(BaseTransform):
    def __call__(self, data: Data) -> Data:
        mask = data.y.new_zeros(data.num_nodes, dtype=bool)
        mask[data.edge_index[0]] = True
        mask[data.edge_index[1]] = True
        data = data.subgraph(mask)
        return data

core/data/transforms/remove_isolated_nodes.py
thupchnsky/dp-gnn
from torch_geometric.utils import remove_self_loops
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data


class RemoveSelfLoops(BaseTransform):
    def __call__(self, data: Data) -> Data:
        if hasattr(data, 'edge_index') and data.edge_index is not None:
            data.edge_index, _ = remove_self_loops(data.edge_index)
        if hasattr(data, 'adj_t'):
            data.adj_t = data.adj_t.remove_diag()
        return data

core/data/transforms/remove_self_loops.py
thupchnsky/dp-gnn
import torch
from torch_geometric.transforms import BaseTransform
from torch_geometric.data import Data


class RandomDataSplit(BaseTransform):
    def __init__(self, num_nodes_per_class, train_ratio=0.7, test_ratio=0.2):
        self.num_nodes_per_class = num_nodes_per_class
        self.train_ratio = train_ratio
        self.test_ratio = test_ratio

    def __call__(self, data: Data) -> Data:
        y = data.y
        num_classes = y.max().item() + 1
        num_train_nodes_per_class = int(self.num_nodes_per_class * self.train_ratio)
        num_test_nodes_per_class = int(self.num_nodes_per_class * self.test_ratio)

        train_mask = torch.zeros_like(y, dtype=torch.bool)
        test_mask = torch.zeros_like(y, dtype=torch.bool)
        val_mask = torch.zeros_like(y, dtype=torch.bool)
        
        for c in range(num_classes):
            idx = (y == c).nonzero(as_tuple=False).view(-1)
            num_nodes = idx.size(0)
            if num_nodes >= self.num_nodes_per_class:    
                idx = idx[torch.randperm(idx.size(0))][:self.num_nodes_per_class]
                train_mask[idx[:num_train_nodes_per_class]] = True
                test_mask[idx[num_train_nodes_per_class:num_train_nodes_per_class + num_test_nodes_per_class]] = True
                val_mask[idx[num_train_nodes_per_class + num_test_nodes_per_class:]] = True

        data.train_mask = train_mask
        data.test_mask = test_mask
        data.val_mask = val_mask

        return data
core/data/transforms/split.py
thupchnsky/dp-gnn
from core.datasets.facebook import Facebook
from core.datasets.amazon import Amazon
from core.datasets.loader import DatasetLoader
core/datasets/__init__.py
thupchnsky/dp-gnn
import torch
from ogb.nodeproppred import PygNodePropPredDataset


class Amazon(PygNodePropPredDataset):
    def __init__(self, root: str, transform=None, pre_transform=None):
        super().__init__(root=root, name='ogbn-products', transform=None, pre_transform=pre_transform)
        self.data.__num_nodes__ = self.data.x.size(0)
        self.data.y = self.data.y.view(-1)
        split_idx = self.get_idx_split()
        for split, idx in split_idx.items():
            mask = torch.zeros(self.data.num_nodes, dtype=bool).scatter(0, idx, True)
            split = 'val' if split == 'valid' else split
            self.data[f'{split}_mask'] = mask
        if transform:
            self.data = transform(self.data)
core/datasets/amazon.py
thupchnsky/dp-gnn
#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
#
# Distributed under terms of the MIT license.

"""
This is a script for contexual SBM model and its dataset generator.
contains functions:
        ContextualSBM
        parameterized_Lambda_and_mu
        save_data_to_pickle
    class:
        dataset_ContextualSBM

"""

import numpy as np
import torch
from torch_geometric.data import Data
import pickle
from datetime import datetime
import os.path as osp
import os
import argparse
from torch_geometric.data import InMemoryDataset


def index_to_mask(index, size):
    mask = torch.zeros(size, dtype=torch.bool, device=index.device)
    mask[index] = 1
    return mask


def random_planetoid_splits(data, num_classes, percls_trn=20, val_lb=500, Flag=0):
    # Set new random planetoid splits:
    # * round(train_rate*len(data)/num_classes) * num_classes labels for training
    # * val_rate*len(data) labels for validation
    # * rest labels for testing

    indices = []
    for i in range(num_classes):
        index = (data.y == i).nonzero().view(-1)
        index = index[torch.randperm(index.size(0))]
        indices.append(index)

    train_index = torch.cat([i[:percls_trn] for i in indices], dim=0)

    if Flag is 0:
        rest_index = torch.cat([i[percls_trn:] for i in indices], dim=0)
        rest_index = rest_index[torch.randperm(rest_index.size(0))]

        data.train_mask = index_to_mask(train_index, size=data.num_nodes)
        data.val_mask = index_to_mask(rest_index[:val_lb], size=data.num_nodes)
        data.test_mask = index_to_mask(
            rest_index[val_lb:], size=data.num_nodes)
    else:
        val_index = torch.cat([i[percls_trn:percls_trn+val_lb]
                               for i in indices], dim=0)
        rest_index = torch.cat([i[percls_trn+val_lb:] for i in indices], dim=0)
        rest_index = rest_index[torch.randperm(rest_index.size(0))]

        data.train_mask = index_to_mask(train_index, size=data.num_nodes)
        data.val_mask = index_to_mask(val_index, size=data.num_nodes)
        data.test_mask = index_to_mask(rest_index, size=data.num_nodes)
    return data


def ContextualSBM(n, d, Lambda, p, mu, train_percent=0.01):
    # n = 800 #number of nodes
    # d = 5 # average degree
    # Lambda = 1 # parameters
    # p = 1000 # feature dim
    # mu = 1 # mean of Gaussian
    gamma = n/p

    c_in = d + np.sqrt(d)*Lambda
    c_out = d - np.sqrt(d)*Lambda
    y = np.ones(n)
    y[int(n/2)+1:] = -1
    y = np.asarray(y, dtype=int)

    # creating edge_index
    edge_index = [[], []]
    for i in range(n-1):
        for j in range(i+1, n):
            if y[i]*y[j] > 0:
                Flip = np.random.binomial(1, c_in/n)
            else:
                Flip = np.random.binomial(1, c_out/n)
            if Flip > 0.5:
                edge_index[0].append(i)
                edge_index[1].append(j)
                edge_index[0].append(j)
                edge_index[1].append(i)

    # creating node features
    x = np.zeros([n, p])
    u = np.random.normal(0, 1/np.sqrt(p), [1, p])
    for i in range(n):
        Z = np.random.normal(0, 1, [1, p])
        x[i] = np.sqrt(mu/n)*y[i]*u + Z/np.sqrt(p)
    data = Data(x=torch.tensor(x, dtype=torch.float32),
                edge_index=torch.tensor(edge_index),
                y=torch.tensor((y + 1) // 2, dtype=torch.int64))
    # order edge list and remove duplicates if any.
    data.coalesce()

    num_class = len(np.unique(y))
    val_lb = int(n * train_percent)
    percls_trn = int(round(train_percent * n / num_class))
    data = random_planetoid_splits(data, num_class, percls_trn, val_lb)

    # add parameters to attribute
    data.Lambda = Lambda
    data.mu = mu
    data.n = n
    data.p = p
    data.d = d
    data.train_percent = train_percent

    return data


def parameterized_Lambda_and_mu(theta, p, n, epsilon=0.1):
    '''
    based on claim 3 in the paper, 

        lambda^2 + mu^2/gamma = 1 + epsilon.

    1/gamma = p/n
    longer axis: 1
    shorter axis: 1/gamma.
    =>
        lambda = sqrt(1 + epsilon) * sin(theta * pi / 2)
        mu = sqrt(gamma * (1 + epsilon)) * cos(theta * pi / 2)
    '''
    from math import pi
    gamma = n / p
    assert (theta >= -1) and (theta <= 1)
    Lambda = np.sqrt(1 + epsilon) * np.sin(theta * pi / 2)
    mu = np.sqrt(gamma * (1 + epsilon)) * np.cos(theta * pi / 2)
    return Lambda, mu


def save_data_to_pickle(data, p2root='../data/', file_name=None):
    '''
    if file name not specified, use time stamp.
    '''
    now = datetime.now()
    surfix = now.strftime('%b_%d_%Y-%H:%M')
    if file_name is None:
        tmp_data_name = '_'.join(['cSBM_data', surfix])
    else:
        tmp_data_name = file_name
    p2cSBM_data = osp.join(p2root, tmp_data_name)
    if not osp.isdir(p2root):
        os.makedirs(p2root)
    with open(p2cSBM_data, 'bw') as f:
        pickle.dump(data, f)
    return p2cSBM_data


class dataset_ContextualSBM(InMemoryDataset):
    r"""Create synthetic dataset based on the contextual SBM from the paper:
    https://arxiv.org/pdf/1807.09596.pdf

    Use the similar class as InMemoryDataset, but not requiring the root folder.

       See `here <https://pytorch-geometric.readthedocs.io/en/latest/notes/
    create_dataset.html#creating-in-memory-datasets>`__ for the accompanying
    tutorial.

    Args:
        root (string): Root directory where the dataset should be saved.
        name (string): The name of the dataset if not specified use time stamp.

        for {n, d, p, Lambda, mu}, with '_' as prefix: intial/feed in argument.
        without '_' as prefix: loaded from data information

        n: number nodes
        d: avg degree of nodes
        p: dimenstion of feature vector.

        Lambda, mu: parameters balancing the mixture of information, 
                    if not specified, use parameterized method to generate.

        epsilon, theta: gap between boundary and chosen ellipsoid. theta is 
                        angle of between the selected parameter and x-axis.
                        choosen between [0, 1] => 0 = 0, 1 = pi/2

        transform (callable, optional): A function/transform that takes in an
            :obj:`torch_geometric.data.Data` object and returns a transformed
            version. The data object will be transformed before every access.
            (default: :obj:`None`)
        pre_transform (callable, optional): A function/transform that takes in
            an :obj:`torch_geometric.data.Data` object and returns a
            transformed version. The data object will be transformed before
            being saved to disk. (default: :obj:`None`)
    """

#     url = 'https://github.com/kimiyoung/planetoid/raw/master/data'

    def __init__(self, root, name=None,
                 n=800, d=5, p=100, Lambda=None, mu=None,
                 epsilon=0.1, theta=0.5,
                 train_percent=0.01,
                 transform=None, pre_transform=None):

        now = datetime.now()
        surfix = now.strftime('%b_%d_%Y-%H:%M')
        if name is None:
            # not specifing the dataset name, create one with time stamp.
            self.name = '_'.join(['cSBM_data', surfix])
        else:
            self.name = name

        self._n = n
        self._d = d
        self._p = p

        self._Lambda = Lambda
        self._mu = mu
        self._epsilon = epsilon
        self._theta = theta
        
        self._train_percent = train_percent

        root = osp.join(root, self.name)
        if not osp.isdir(root):
            os.makedirs(root)
        super(dataset_ContextualSBM, self).__init__(
            root, transform, pre_transform)

#         ipdb.set_trace()
        self.data, self.slices = torch.load(self.processed_paths[0])
        # overwrite the dataset attribute n, p, d, Lambda, mu
        self.Lambda = self.data.Lambda.item()
        self.mu = self.data.mu.item()
        self.n = self.data.n
        self.p = self.data.p
        self.d = self.data.d
        self.train_percent = self.data.train_percent

#     @property
#     def raw_dir(self):
#         return osp.join(self.root, self.name, 'raw')

#     @property
#     def processed_dir(self):
#         return osp.join(self.root, self.name, 'processed')

    @property
    def raw_file_names(self):
        file_names = [self.name]
        return file_names

    @property
    def processed_file_names(self):
        return ['data.pt']

    def download(self):
        for name in self.raw_file_names:
            p2f = osp.join(self.raw_dir, name)
            if not osp.isfile(p2f):
                # file not exist, so we create it and save it there.
                if self._Lambda is None or self._mu is None:
                    # auto generate the lambda and mu parameter by angle theta.
                    self._Lambda, self._mu = parameterized_Lambda_and_mu(self._theta,
                                                                         self._p,
                                                                         self._n,
                                                                         self._epsilon)
                tmp_data = ContextualSBM(self._n,
                                         self._d,
                                         self._Lambda,
                                         self._p,
                                         self._mu,
                                         self._train_percent)

                _ = save_data_to_pickle(tmp_data,
                                        p2root=self.raw_dir,
                                        file_name=self.name)
            else:
                # file exists already. Do nothing.
                pass

    def process(self):
        p2f = osp.join(self.raw_dir, self.name)
        with open(p2f, 'rb') as f:
            data = pickle.load(f)
        data = data if self.pre_transform is None else self.pre_transform(data)
        torch.save(self.collate([data]), self.processed_paths[0])

    def __repr__(self):
        return '{}()'.format(self.name)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--phi', type=float, default=1)
    parser.add_argument('--epsilon', type = float , default = 3.25)
    parser.add_argument('--root', default = '../data/')
    parser.add_argument('--name', default = 'cSBM_demo')
    parser.add_argument('--num_nodes', type = int, default = 800)
    parser.add_argument('--num_features', type = int, default = 1000)
    parser.add_argument('--avg_degree', type = float, default = 5)

    args = parser.parse_args()

    dataset_ContextualSBM(root = args.root,
            name = args.name,
            theta = args.phi,
            epsilon = args.epsilon, 
            n = args.num_nodes,
            d = args.avg_degree,
            p = args.num_features)

core/datasets/cSBM_dataset.py
thupchnsky/dp-gnn
import os
import ssl
import torch
import pandas as pd
from scipy.io import loadmat
from sklearn.preprocessing import LabelEncoder
from torch_geometric.utils import subgraph, from_scipy_sparse_matrix
from torch_geometric.data import Data, InMemoryDataset, download_url


class Facebook(InMemoryDataset):
    url = 'https://github.com/sisaman/pyg-datasets/raw/main/datasets/facebook100/'
    targets = ['status', 'gender', 'major', 'minor', 'housing', 'year']
    available_datasets = [
        "Villanova62", "UCLA26", "Tennessee95", "NYU9", "Carnegie49", "GWU54", "USF51", "Vanderbilt48", 
        "USC35", "Rutgers89", "UConn91", "MIT8", "USFCA72", "UChicago30", "UIllinois20", "UC61", "Cal65", 
        "Yale4", "Northeastern19", "Dartmouth6", "Vermont70", "Northwestern25", "William77", "Harvard1", 
        "Princeton12", "UC64", "Middlebury45", "Haverford76", "Bingham82", "UNC28", "Berkeley13", "Rochester38", 
        "Swarthmore42", "Virginia63", "WashU32", "Columbia2", "NotreDame57", "Bucknell39", "UVA16", "Maine59", 
        "MU78", "Simmons81", "MSU24", "Colgate88", "Temple83", "Cornell5", "Indiana69", "Oklahoma97", "Michigan23", 
        "BU10", "Brown11", "Auburn71", "FSU53", "UGA50", "UCF52", "Howard90", "UCSD34", "Vassar85", "Tufts18", 
        "UPenn7", "Baylor93", "UMass92", "Bowdoin47", "Maryland58", "Penn94", "Wesleyan43", "UC33", 
        "Rice31", "UCSC68", "Smith60", "Caltech36", "Hamilton46", "Oberlin44", "American75", "Mich67", 
        "Mississippi66", "Williams40", "UCSB37", "Amherst41", "Duke14", "Pepperdine86", "Wake73", "Lehigh96", 
        "Reed98", "Tulane29", "Texas84", "Wellesley22", "JMU79", "Santa74", "Wisconsin87", "Stanford3", 
        "Texas80", "UF21", "JohnsHopkins55", "Syracuse56", "BC17", "Georgetown15", "Trinity100", "Brandeis99", "Emory27"
    ]

    def __init__(self, root: str, name: str, target='status', transform=None, pre_transform=None):
        self.name = name
        self.target = target
        super().__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_dir(self) -> str:
        return os.path.join(self.root, self.name, 'raw')

    @property
    def raw_file_names(self) -> str:
        return self.name + '.mat'

    @property
    def processed_dir(self) -> str:
        return os.path.join(self.root, self.name, 'processed')

    @property
    def processed_file_names(self) -> str:
        return f'data-{self.name}-{self.target}.pt'

    def download(self):
        context = ssl._create_default_https_context
        ssl._create_default_https_context = ssl._create_unverified_context
        download_url(f'{self.url}/{self.raw_file_names}', self.raw_dir)
        ssl._create_default_https_context = context

    def process(self):
        mat = loadmat(os.path.join(self.raw_dir, self.raw_file_names))
        features = pd.DataFrame(mat['local_info'][:, :-1], columns=self.targets)
        y = torch.from_numpy(LabelEncoder().fit_transform(features[self.target]))
        if 0 in features[self.target].values:
            y = y - 1

        x = features.drop(columns=self.target).replace({0: None})
        x = torch.tensor(pd.get_dummies(x).values, dtype=torch.float)
        edge_index = from_scipy_sparse_matrix(mat['A'])[0]

        # removed unlabled nodes
        subset = y >= 0
        edge_index, _ = subgraph(subset, edge_index, relabel_nodes=True, num_nodes=len(y))
        x = x[subset]
        y = y[subset]

        data = Data(x=x, edge_index=edge_index, y=y)

        if self.pre_transform is not None:
            data = self.pre_transform(data)

        torch.save(self.collate([data]), self.processed_paths[0])

    def __repr__(self) -> str:
        return f'Facebook100-{self.name}()'


core/datasets/facebook.py
thupchnsky/dp-gnn
import os
from functools import partial
from typing import Annotated
import torch
from core import console
from torch_geometric.data import Data
from torch_geometric.datasets import Planetoid, WikipediaNetwork
from torch_geometric.datasets import Amazon as Amazon_pyg
from torch_geometric.transforms import Compose, ToSparseTensor, RandomNodeSplit
from core.args.utils import ArgInfo
from core.data.transforms import FilterClassByCount
from core.data.transforms import RemoveSelfLoops
from core.data.transforms import RemoveIsolatedNodes
from core.datasets import Facebook
from core.datasets import Amazon
from core.datasets.cSBM_dataset import dataset_ContextualSBM
from core.utils import dict2table


class DatasetLoader:
    supported_datasets = {
        'amazon': partial(Amazon, 
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=100000, remove_unlabeled=True)
            ])
        ),
        'chameleon': partial(WikipediaNetwork,name="chameleon",
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        ),
        'squirrel': partial(WikipediaNetwork,name="squirrel",
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        ),
        'photo': partial(Amazon_pyg,name = 'photo',
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        ),
        'computers': partial(Amazon_pyg,name = 'computers',
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        ),
        'facebook': partial(Facebook, name='UIllinois20', target='year', 
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=1000, remove_unlabeled=True)
            ])
        ),
        'cora': partial(Planetoid, name='Cora',
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        ),
        'pubmed': partial(Planetoid, name='PubMed',
            transform=Compose([
                RandomNodeSplit(num_val=0.1, num_test=0.15), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        ),
        'cSBM': partial(dataset_ContextualSBM, epsilon=3.25, n=30000, d=10, p=2000,
            transform=Compose([
                RandomNodeSplit(num_val=0.025, num_test=0.95), 
                FilterClassByCount(min_count=50, remove_unlabeled=True)
            ])
        )
    }

    def __init__(self,
                 dataset:    Annotated[str, ArgInfo(help='name of the dataset', choices=supported_datasets)] = 'facebook',
                 data_dir:   Annotated[str, ArgInfo(help='directory to store the dataset')] = 'exp_datasets',
                 ):

        self.name = dataset
        self.data_dir = data_dir

    def load(self, verbose=False, csbm_phi=0.0) -> Data:
        if self.name == 'cSBM':
            data = self.supported_datasets[self.name](root=os.path.join(self.data_dir, self.name), name=f'cSBM_data_phi_{csbm_phi}', theta=csbm_phi)[0]
        else:
            data = self.supported_datasets[self.name](root=os.path.join(self.data_dir, self.name))[0]
        data = Compose([RemoveSelfLoops(), RemoveIsolatedNodes(), ToSparseTensor()])(data)

        if verbose:
            self.print_stats(data)

        return data

    def print_stats(self, data: Data):
        nodes_degree: torch.Tensor = data.adj_t.sum(dim=1)
        baseline: float = (data.y[data.test_mask].unique(return_counts=True)[1].max().item() * 100 / data.test_mask.sum().item())
        train_ratio: float = data.train_mask.sum().item() / data.num_nodes * 100
        val_ratio: float = data.val_mask.sum().item() / data.num_nodes * 100
        test_ratio: float = data.test_mask.sum().item() / data.num_nodes * 100

        stat = {
            'nodes': f'{data.num_nodes:,}',
            'edges': f'{data.num_edges:,}',
            'features': f'{data.num_features:,}',
            'classes': f'{int(data.y.max() + 1)}',
            'mean degree': f'{nodes_degree.mean():.2f}',
            'median degree': f'{nodes_degree.median()}',
            'train/val/test (%)': f'{train_ratio:.1f}/{val_ratio:.1f}/{test_ratio:.1f}',
            'baseline acc (%)': f'{baseline:.2f}'
        }

        table = dict2table(stat, num_cols=2, title=f'dataset: [yellow]{self.name}[/yellow]')
        console.info(table)
        console.print()

core/datasets/loader.py
thupchnsky/dp-gnn
from core.loggers.base import LoggerBase
from core.loggers.csv import CSVLogger
from core.loggers.wandb import WandbLogger
from core.loggers.factory import Logger


core/loggers/__init__.py
thupchnsky/dp-gnn
import functools
from typing import Annotated, Callable
from uuid import uuid1
from abc import ABC, abstractmethod
from torch.nn import Module
from core.args.utils import ArgInfo
from core.utils import RT


def if_enabled(func: Callable[..., RT]) -> Callable[..., RT]:
    @functools.wraps(func)
    def wrapper(self: LoggerBase, *args, **kwargs) -> RT:
        if self.enabled:
            return func(self, *args, **kwargs)
    return wrapper


class LoggerBase(ABC):
    def __init__(self, 
                 project:       Annotated[str,  ArgInfo(help="project name for logger")] = 'GAP',
                 output_dir:    Annotated[str,  ArgInfo(help="directory to store the results")] = './output',
                 experiment_id: str=str(uuid1()), 
                 enabled:       bool=True, 
                 config:        dict={}
                 ):
        self.project = project
        self.experiment_id = experiment_id
        self.output_dir = output_dir
        self.enabled = enabled
        self.config = config

    @property
    @abstractmethod
    def experiment(self): pass

    @abstractmethod
    def log(self, metrics: dict[str, object]): pass
    
    @abstractmethod
    def log_summary(self, metrics: dict[str, object]): pass

    @abstractmethod
    def watch(self, model: Module, **kwargs): pass
    
    @abstractmethod
    def finish(self): pass

    def enable(self): 
        self.enabled = True

    def disable(self): 
        self.enabled = False

core/loggers/base.py
thupchnsky/dp-gnn
import os
import pandas as pd
from torch.nn import Module
from core.loggers.base import LoggerBase, if_enabled


class CSVLogger(LoggerBase):
    @property
    def experiment(self) -> dict[str, object]:
        if not hasattr(self, '_experiment'):
            self._experiment = self.config
        return self._experiment

    @if_enabled
    def log(self, metrics: dict[str, object]):
        self.experiment.update(metrics)

    @if_enabled
    def log_summary(self, metrics: dict[str, object]):
        self.experiment.update(metrics)

    def watch(self, model: Module, **kwargs):
        pass

    @if_enabled
    def finish(self):
        log_dir = os.path.join(self.output_dir, 'csv', self.project)
        os.makedirs(log_dir, exist_ok=True)
        self.experiment.pop("priv_weights_list", None)
        df = pd.DataFrame(self.experiment, index=[0])
        df.to_csv(os.path.join(log_dir, f'{self.experiment_id}.csv'), index=False)

core/loggers/csv.py
thupchnsky/dp-gnn
from typing import Annotated
from core import console
from core import globals
from core.args.utils import ArgInfo
from core.loggers.base import LoggerBase
from core.loggers.csv import CSVLogger
from core.loggers.wandb import WandbLogger


class Logger:
    supported_loggers = {
        'csv': CSVLogger,
        'wandb': WandbLogger,
    }

    _instance: LoggerBase = None
    
    @classmethod
    def get_instance(cls) -> LoggerBase:
        return cls._instance

    def __new__(cls) -> LoggerBase:
        if cls._instance is None:
            raise RuntimeError('Logger is not initialized, call Logger.setup() first')
        return cls._instance

    @classmethod
    def setup(cls,
        logger:   Annotated[str,  ArgInfo(help='select logger type', choices=supported_loggers)] = 'csv',
        **kwargs: Annotated[dict, ArgInfo(help='additional kwargs for the underlying logger', bases=[LoggerBase])],
        ) -> LoggerBase:
        
        if globals['debug']:
            logger = 'wandb'
            kwargs['enabled'] = True
            kwargs['project'] += '-DEBUG'
            console.debug(f'debug mode: wandb logger is enabled for project {kwargs["project"]}')
        
        LoggerCls = cls.supported_loggers[logger]
        cls._instance = LoggerCls(**kwargs)
        return cls._instance

core/loggers/factory.py
thupchnsky/dp-gnn
import os

from torch.nn import Module
from core.loggers.base import LoggerBase, if_enabled

try: 
    import wandb
    from wandb.wandb_run import Run
except ImportError: 
    wandb = None


class WandbLogger(LoggerBase):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        if wandb is None:
            raise ImportError('wandb is not installed yet, install it with `pip install wandb`.')

    @property
    def experiment(self):
        if not hasattr(self, '_experiment'):
            os.environ["WANDB_SILENT"] = "true"
            settings = wandb.Settings(start_method="fork")
            os.makedirs(self.output_dir, exist_ok=True)

            self._experiment = wandb.init(
                project=self.project,
                dir=self.output_dir,
                reinit=True, 
                resume='allow', 
                config=self.config, 
                save_code=True,
                settings=settings
            )

        return self._experiment

    @if_enabled
    def log(self, metrics: dict[str, object]):
        self.experiment.log(metrics)

    @if_enabled
    def log_summary(self, metrics: dict[str, object]):
        for metric, value in metrics.items():
            self.experiment.summary[metric] = value

    @if_enabled
    def watch(self, model: Module, **kwargs):
        self.experiment.watch(model, **kwargs)

    @if_enabled
    def finish(self):
        self.experiment.finish()

core/loggers/wandb.py
thupchnsky/dp-gnn
from abc import ABC, abstractmethod
from typing import Optional
from torch import Tensor
from torch_geometric.data import Data
from core.modules.base import Metrics


class MethodBase(ABC):

    @abstractmethod
    def fit(self, data: Data, prefix: str = '') -> Metrics:
        """Fit the model to the given data."""

    @abstractmethod
    def test(self, data: Optional[Data] = None, prefix: str = '') -> Metrics:
        """Test the model on the given data, or the training data if data is None."""

    @abstractmethod
    def predict(self, data: Optional[Data] = None) -> Tensor:
        """Predict the labels for the given data, or the training data if data is None."""

core/methods/base.py
thupchnsky/dp-gnn
from core.methods.node.base import NodeClassification
from core.methods.node.gap.gap_inf import GAP
from core.methods.node.gap.gap_edp import EdgePrivGAP
from core.methods.node.gap.gap_ndp import NodePrivGAP
from core.methods.node.sage.sage_inf import SAGE
from core.methods.node.sage.sage_edp import EdgePrivSAGE
from core.methods.node.sage.sage_ndp import NodePrivSAGE
from core.methods.node.mlp.mlp import MLP
from core.methods.node.mlp.mlp_dp import PrivMLP
from core.methods.node.dpdgc.dpdgc_inf import DPDGC
from core.methods.node.dpdgc.dpdgc_edp import EdgePrivDPDGC
from core.methods.node.dpdgc.dpdgc_ndp import NodePrivDPDGC
from core.methods.node.dpdgc.dpdgc_kndp import KNeighborPrivDPDGC

supported_methods = {
    'gap-inf':  GAP,
    'gap-edp':  EdgePrivGAP,
    'gap-ndp':  NodePrivGAP,
    'sage-inf': SAGE,
    'sage-edp': EdgePrivSAGE,
    'sage-ndp': NodePrivSAGE,
    'mlp':      MLP,
    'mlp-dp':   PrivMLP,
    'dpdgc-inf': DPDGC,
    'dpdgc-edp': EdgePrivDPDGC,
    'dpdgc-ndp': NodePrivDPDGC,
    'dpdgc-kndp': KNeighborPrivDPDGC
}

core/methods/node/__init__.py
thupchnsky/dp-gnn
from abc import abstractmethod
from typing import Annotated, Literal, Optional, Union
import torch
from torch import Tensor
from torch.optim import Adam, SGD, Optimizer
from torch_geometric.data import Data
from core.args.utils import ArgInfo
from core.data.loader import NodeDataLoader, SparseAdjRowLoader
from core import globals
from core.methods.base import MethodBase
from core.modules.base import TrainableModule, Metrics, Stage
from core import console
from core.trainer import Trainer


class NodeClassification(MethodBase):
    def __init__(self, 
                 num_classes:     int, 
                 epochs:          Annotated[int,   ArgInfo(help='number of epochs for training')] = 100,
                 optimizer:       Annotated[str,   ArgInfo(help='optimization algorithm', choices=['sgd', 'adam'])] = 'adam',
                 learning_rate:   Annotated[float, ArgInfo(help='learning rate', option='--lr')] = 0.01,
                 weight_decay:    Annotated[float, ArgInfo(help='weight decay (L2 penalty)')] = 0.0,
                 batch_size:      Annotated[Union[Literal['full'], int],   
                                                   ArgInfo(help='batch size, or "full" for full-batch training')] = 'full',
                 full_batch_eval: Annotated[bool,  ArgInfo(help='if true, then model uses full-batch evaluation')] = True,
                 device:          Annotated[str,   ArgInfo(help='device to use', choices=['cpu', 'cuda', 'cuda:0', 'cuda:1'])] = 'cuda',
                 **trainer_args:  Annotated[dict,  ArgInfo(help='extra options passed to the trainer class', bases=[Trainer])]
                 ):

        self.num_classes = num_classes
        self.epochs = epochs
        self.optimizer_name = optimizer
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.batch_size = batch_size
        self.full_batch_eval = full_batch_eval
        self.device = device

        if self.device == 'cuda' and not torch.cuda.is_available():
            console.warning('CUDA is not available, proceeding with CPU') 
            self.device = 'cpu'

        self.data = None  # data is kept for caching purposes
        self.trainer = self.configure_trainer(**trainer_args)

    @property
    @abstractmethod
    def classifier(self) -> TrainableModule:
        """Return the underlying classifier."""

    def reset_parameters(self):
        self.classifier.reset_parameters()
        self.trainer.reset()
        self.data = None

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '') -> Metrics:
        """Fit the model to the given data."""
        if learn == 'transductive':
            self.data = data.to(self.device, non_blocking=True)
            train_metrics = self._train(self.data, prefix=prefix)
            # test uses cached results
            test_metrics = self.test(self.data, prefix=prefix)
        else:
            # data and self.data are already different now
            train_metrics = self._train(data, prefix=prefix)
            test_metrics = self.test(self.data, prefix=prefix)
        return {**train_metrics, **test_metrics}

    def test(self, data: Optional[Data] = None, prefix: str = '') -> Metrics:
        """Predict the labels for the given data, or the training data if data is None."""
        if data is None:
            data = self.data
        
        data = data.to(self.device, non_blocking=True)

        test_metics = self.trainer.test(
            dataloader=self.data_loader(data, 'test'),
            prefix=prefix,
        )

        return test_metics

    def predict(self, data: Optional[Data] = None) -> Tensor:
        """Predict the labels for the given data, or the training data if data is None."""
        if data is None:
            data = self.data

        data = data.to(self.device, non_blocking=True)
        return self.classifier.predict(data)

    def _train(self, data: Data, prefix: str = '') -> Metrics:
        console.info('training classifier')
        self.classifier.to(self.device)

        metrics = self.trainer.fit(
            model=self.classifier,
            epochs=self.epochs,
            optimizer=self.configure_optimizer(),
            train_dataloader=self.data_loader(data, 'train'), 
            val_dataloader=self.data_loader(data, 'val'),
            test_dataloader=self.data_loader(data, 'test') if globals['debug'] else None,
            checkpoint=True,
            prefix=prefix,
        )

        return metrics

    def configure_trainer(self, **kwargs) -> Trainer:
        trainer = Trainer(
            monitor='val/acc', 
            monitor_mode='max', 
            **kwargs
        )
        return trainer

    def configure_optimizer(self) -> Optimizer:
        Optim = {'sgd': SGD, 'adam': Adam}[self.optimizer_name]
        return Optim(self.classifier.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)

    def data_loader(self, data: Data, stage: Stage) -> NodeDataLoader:
        """Return a dataloader for the given stage."""
        
        batch_size = 'full' if (stage != 'train' and self.full_batch_eval) else self.batch_size
        dataloader = NodeDataLoader(
            data=data, 
            stage=stage,
            batch_size=batch_size, 
            shuffle=(stage == 'train'), 
            drop_last=True,
            poisson_sampling=False,
        )

        return dataloader
    
    def adj_data_loader(self, data: Data, stage: Stage) -> SparseAdjRowLoader:
        """Return a dataloader for the given stage."""
        
        batch_size = 'full' if (stage != 'train' and self.full_batch_eval) else self.batch_size
        dataloader = SparseAdjRowLoader(
            data=data,
            stage=stage,
            batch_size=batch_size,
            shuffle=(stage == 'train'),
            drop_last=False,
            poisson_sampling=False,
        )

        return dataloader

core/methods/node/base.py
thupchnsky/dp-gnn
import numpy as np
import torch
from typing import Annotated, Literal, Union
from torch_geometric.data import Data
from torch_sparse import SparseTensor, matmul
from opacus.optimizers import DPOptimizer
from core import console
from core.args.utils import ArgInfo
from core.data.loader import SparseAdjRowLoader
from core.methods.node import DPDGC
from core.privacy.mechanisms import ComposedNoisyMechanism
from core.privacy.algorithms import PMA, NoisySGD
from core.data.transforms import BoundOutDegree
from core.modules.base import Metrics, Stage


class EdgePrivDPDGC (DPDGC):
    """edge-private DPDGC method"""

    def __init__(self,
                 num_features,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 max_grad_norm: Annotated[float, ArgInfo(help='maximum norm of the per-sample gradients')] = 1.0,
                 batch_size:    Annotated[int,   ArgInfo(help='batch size')] = 256,
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[DPDGC], exclude=['encoder_batch_norm'])]
                 ):

        super().__init__(num_features=num_features,
                         num_classes=num_classes,
                         encoder_batch_norm=False,
                         batch_size=batch_size, 
                         **kwargs)
        self.epsilon = epsilon
        self.delta = delta
        self.max_grad_norm = max_grad_norm
        self.num_train_nodes = None  # will be used to set delta if it is 'auto'

    def calibrate(self):
        self.pma_mechanism = PMA(noise_scale=0.0, hops=1)

        self.adj_encoder_noisy_sgd = NoisySGD(
            noise_scale=0.0,
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size,
            epochs=self.encoder_epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True,
        )

        composed_mechanism = ComposedNoisyMechanism(
            noise_scale=0.0,
            mechanism_list=[
                self.adj_encoder_noisy_sgd,
                self.pma_mechanism,
            ]
        )

        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_edges)))
                console.info('delta = %.0e' % delta)
            
            self.noise_scale = composed_mechanism.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')

        self._adj_encoder = self.adj_encoder_noisy_sgd.prepare_module(self._adj_encoder)

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '', p: float = 0.6) -> Metrics:
        num_train_nodes = data.train_mask.sum().item()

        if num_train_nodes != self.num_train_nodes:
            self.num_train_nodes = num_train_nodes
            self.num_edges = data.num_edges
            self.calibrate() 

        return super().fit(data, learn=learn, prefix=prefix)

    def perturb(self, x) -> torch.Tensor:
        x = self.pma_mechanism(x, sensitivity=self.norm_scale)
        return x

    def configure_adj_encoder_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_adj_encoder_optimizer()
        optimizer = self.adj_encoder_noisy_sgd.prepare_optimizer(optimizer, 2.0) # *2 for replacement DP
        return optimizer
    
    def adj_data_loader(self,data: Data, stage: Stage) -> SparseAdjRowLoader:
        dataloader = super().adj_data_loader(data,stage)
        if stage == 'train':
            dataloader.wor_sampling = True
            # dataloader.poisson_sampling = True
        return dataloader

core/methods/node/dpdgc/dpdgc_edp.py
thupchnsky/dp-gnn
import torch
from typing import Annotated, Optional
import torch.nn.functional as F
from torch.optim import Adam, SGD, Optimizer
from torch_geometric.data import Data
from torch_geometric.utils import subgraph
from torch_geometric.transforms import ToSparseTensor
from torch_sparse import SparseTensor, matmul
from core import console
from core.args.utils import ArgInfo
from core.methods.node.base import NodeClassification
from core.models.multi_mlp import MultiMLP
from core.modules.base import Metrics
from core.modules.node.cm import DPDGCClassificationModule
from core.modules.node.em import AdjEncoderModule
from core.data.transforms import BoundOutDegree, BoundDegree, AddSelfLoops, NormalizeDegree


class DPDGC (NodeClassification):
    """Non-private DGC method"""

    supported_activations = {
        'relu': torch.relu_,
        'selu': torch.selu_,
        'tanh': torch.tanh,
    }

    def __init__(self,
                 num_features,
                 num_classes,
                 hidden_dim:           Annotated[int,   ArgInfo(help='hidden dimension of the MLP layers')] = 64,
                 encoder_head_layers:  Annotated[int,   ArgInfo(help='number of MLP layers in MLP encoder')] = 1,
                 base_layers:          Annotated[int,   ArgInfo(help='number of base MLP layers in classifier')] = 2,
                 head_layers:          Annotated[int,   ArgInfo(help='number of head MLP layers in classifier')] = 1,
                 combine:              Annotated[str,   ArgInfo(help='combination type of transformed hops', choices=MultiMLP.supported_combinations)] = 'cat',
                 activation:           Annotated[str,   ArgInfo(help='type of activation function', choices=supported_activations)] = 'selu',
                 dropout:              Annotated[float, ArgInfo(help='general dropout rate')] = 0.5,
                 batch_norm:           Annotated[bool,  ArgInfo(help='if true, then model uses batch normalization')] = True,
                 encoder_batch_norm:   Annotated[bool,  ArgInfo(help='if true, then encoder module uses batch normalization')] = True,
                 encoder_epochs:       Annotated[int,   ArgInfo(help='number of epochs for MLP encoder pre-training (ignored if encoder_layers=0)')] = 50,
                 encoder_lr:           Annotated[float,   ArgInfo(help='learning rate for GRAPH encoder module')] = 5e-3,
                 encoder_wd:           Annotated[float,   ArgInfo(help='weight decay for GRAPH encoder module')] = 1e-5,
                 encoder_dropout:      Annotated[float,   ArgInfo(help='dropout rate for GRAPH encoder module')] = 0.0,
                 norm_scale:           Annotated[float,   ArgInfo(help='scale of MLP_A row norm')] = 1e-8,
                 **kwargs:             Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[NodeClassification])]
                 ):

        super().__init__(num_classes, **kwargs)

        # adj encoder specific learning parameters
        self.encoder_head_layers = encoder_head_layers
        self.encoder_epochs = encoder_epochs
        self.encoder_lr = encoder_lr
        self.encoder_wd = encoder_wd
        self.encoder_dropout = encoder_dropout
        self.norm_scale = norm_scale
        
        activation_fn = self.supported_activations[activation]
        
        self._adj_encoder = AdjEncoderModule(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            head_layers=encoder_head_layers,
            dropout=encoder_dropout,
            activation_fn=activation_fn,
            batch_norm=encoder_batch_norm,
            norm_scale=norm_scale,
        )
        
        self._classifier = DPDGCClassificationModule(
            num_features=num_features,
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            base_layers=base_layers,
            head_layers=head_layers,
            combination=combine,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
        )

    @property
    def classifier(self) -> DPDGCClassificationModule:
        return self._classifier

    def reset_parameters(self):
        self._adj_encoder.reset_parameters()
        super().reset_parameters()

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '') -> Metrics:
        self.learning_setting = learn
        assert learn == 'transductive', 'DGC only support transductive setting for now!'
        
        # in transductive setting we will need to cache the computed data.x for test and predict purpose
        self.data = data.to(self.device, non_blocking=True)
        xA, adj_pretrain_metrics = self.pretrain_adj_encoder(self.data, prefix=prefix)
        xA = self.perturb(xA)
        console.info('adj encoder pretrain statistics')
        console.info(adj_pretrain_metrics)
        
        # concatenate the features for later use
        self.data.x = torch.cat([self.data.x, xA], dim=-1)

        return super().fit(self.data, learn=learn, prefix=prefix)  # train classifier
        
    def test(self, data: Optional[Data] = None, prefix: str = '') -> Metrics:
        assert self.learning_setting == 'transductive', 'DGC only support transductive setting for now!'
        if (data is None or data == self.data) and self.learning_setting == 'transductive':
            # only in transductive we can use cached results
            data = self.data
        return super().test(data, prefix=prefix)
    
    def predict(self, data: Optional[Data] = None) -> torch.Tensor:
        assert self.learning_setting == 'transductive', 'DGC only support transductive setting for now!'
        if (data is None or data == self.data) and self.learning_setting == 'transductive':
            data = self.data
        return super().predict(data)
    
    def perturb(self, x) -> torch.Tensor:
        # do not add pma noise for non-dp methods
        return x

    def pretrain_adj_encoder(self, data: Data, prefix: str) -> Data:
        console.info('pretraining adj encoder (MLP_A)')
        self._adj_encoder.to(self.device)
        
        best_metrics = self.trainer.fit(
            model=self._adj_encoder,
            epochs=self.encoder_epochs,
            optimizer=self.configure_adj_encoder_optimizer(),
            train_dataloader=self.adj_data_loader(data, 'train'), 
            val_dataloader=self.adj_data_loader(data, 'val'),
            test_dataloader=None,
            checkpoint=True,
            prefix=f'{prefix}encoder/',
        )
        
        self.trainer.reset()
        xA = self._adj_encoder.predict(data)
        return xA, best_metrics
    
    def configure_adj_encoder_optimizer(self) -> Optimizer:
        Optim = {'sgd': SGD, 'adam': Adam}[self.optimizer_name]
        return Optim(self._adj_encoder.parameters(), lr=self.encoder_lr, weight_decay=self.encoder_wd)

core/methods/node/dpdgc/dpdgc_inf.py
thupchnsky/dp-gnn
import numpy as np
import torch
from typing import Annotated, Literal, Union
from torch_geometric.data import Data
from torch_sparse import SparseTensor, matmul
from opacus.optimizers import DPOptimizer
from core import console
from core.args.utils import ArgInfo
from core.data.loader import SparseAdjRowLoader, NodeDataLoader
from core.methods.node import DPDGC
from core.privacy.mechanisms import ComposedNoisyMechanism
from core.privacy.algorithms import PMA, NoisySGD
from core.data.transforms import BoundOutDegree
from core.modules.base import Metrics, Stage


class KNeighborPrivDPDGC (DPDGC):
    """k-neighbor-level-private DPDGC method"""

    def __init__(self,
                 num_features,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 nk:            Annotated[int,   ArgInfo(help='max number of edges to change in adjacent dataset')] = 1,
                 max_grad_norm: Annotated[float, ArgInfo(help='maximum norm of the per-sample gradients')] = 1.0,
                 batch_size:    Annotated[int,   ArgInfo(help='batch size')] = 256,
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[DPDGC], exclude=['encoder_batch_norm', 'batch_norm'])]
                 ):

        super().__init__(num_features=num_features,
                         num_classes=num_classes,
                         encoder_batch_norm=False,
                         batch_norm=False,
                         batch_size=batch_size, 
                         **kwargs)
        self.epsilon = epsilon
        self.delta = delta
        self.max_grad_norm = max_grad_norm
        self.num_train_nodes = None  # will be used to set delta if it is 'auto'
        self.nk = nk
        
    def calibrate(self):
        self.pma_mechanism = PMA(noise_scale=0.0, hops=1)

        self.adj_encoder_noisy_sgd = NoisySGD(
            noise_scale=0.0,
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size,
            epochs=self.encoder_epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )

        self.classifier_noisy_sgd = NoisySGD(
            noise_scale=0.0,
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size,
            epochs=self.epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )
        
        if self.nk > 0:
            composed_mechanism = ComposedNoisyMechanism(
                noise_scale=0.0,
                mechanism_list=[
                    self.adj_encoder_noisy_sgd,
                    self.pma_mechanism,
                    self.classifier_noisy_sgd,
                ]
            )
        else:
            composed_mechanism = ComposedNoisyMechanism(
                noise_scale=0.0,
                mechanism_list=[
                    self.adj_encoder_noisy_sgd,
                    self.classifier_noisy_sgd,
                ]
            )

        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_train_nodes))) 
                console.info('delta = %.0e' % delta)
            
            self.noise_scale = composed_mechanism.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')
        
        self._adj_encoder = self.adj_encoder_noisy_sgd.prepare_module(self._adj_encoder)
        self._classifier = self.classifier_noisy_sgd.prepare_module(self._classifier)

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '', p: float = 0.6) -> Metrics:
        num_train_nodes = data.train_mask.sum().item()

        if num_train_nodes != self.num_train_nodes:
            self.num_train_nodes = num_train_nodes
            self.calibrate() 
    
        return super().fit(data, learn=learn, prefix=prefix)

    def perturb(self, x) -> torch.Tensor:
        # Note that the node to be replaced will be take care of by the following DPSGD. 
        # Hence, the effective sensitivity is sqrt(nk)
        if self.nk > 0:
            x = self.pma_mechanism(x, sensitivity=self.norm_scale*np.sqrt(self.nk))
        return x

    def configure_adj_encoder_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_adj_encoder_optimizer()
        # noise_factor = 2 for replacement dp.
        optimizer = self.adj_encoder_noisy_sgd.prepare_optimizer(optimizer, 2.0*np.sqrt(self.nk+1))
        return optimizer
    
    def configure_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_optimizer()
        optimizer = self.classifier_noisy_sgd.prepare_optimizer(optimizer, 2.0) # replacement DP
        return optimizer
    
    def adj_data_loader(self,data: Data, stage: Stage) -> SparseAdjRowLoader:
        dataloader = super().adj_data_loader(data,stage)
        if stage == 'train':
            dataloader.wor_sampling = True
        return dataloader
    
    def data_loader(self, data: Data, stage: Stage) -> NodeDataLoader:
        dataloader = super().data_loader(data, stage)
        if stage == 'train':
            dataloader.wor_sampling = True
        return dataloader

core/methods/node/dpdgc/dpdgc_kndp.py
thupchnsky/dp-gnn
import numpy as np
import torch
from typing import Annotated, Literal, Union
from torch_geometric.data import Data
from torch_sparse import SparseTensor, matmul
from opacus.optimizers import DPOptimizer
from core import console
from core.args.utils import ArgInfo
from core.data.loader import SparseAdjRowLoader, NodeDataLoader
from core.methods.node import DPDGC
from core.privacy.mechanisms import ComposedNoisyMechanism
from core.privacy.algorithms import PMA, NoisySGD
from core.data.transforms import BoundOutDegree
from core.modules.base import Metrics, Stage


class NodePrivDPDGC (DPDGC):
    """Node-private DPDGC method"""

    def __init__(self,
                 num_features,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 max_degree:    Annotated[int,   ArgInfo(help='max degree to sample per each node')] = 100,
                 max_grad_norm: Annotated[float, ArgInfo(help='maximum norm of the per-sample gradients')] = 1.0,
                 batch_size:    Annotated[int,   ArgInfo(help='batch size')] = 256,
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[DPDGC], exclude=['encoder_batch_norm', 'batch_norm'])]
                 ):

        super().__init__(num_features=num_features,
                         num_classes=num_classes,
                         encoder_batch_norm=False,
                         batch_norm=False,
                         batch_size=batch_size, 
                         **kwargs)
        self.epsilon = epsilon
        self.delta = delta
        self.max_grad_norm = max_grad_norm
        self.max_degree = max_degree
        self.num_train_nodes = None  # will be used to set delta if it is 'auto'

    def calibrate(self):
        self.pma_mechanism = PMA(noise_scale=0.0, hops=1)

        self.adj_encoder_noisy_sgd = NoisySGD(
            noise_scale=0.0,
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size,
            epochs=self.encoder_epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )
        self.classifier_noisy_sgd = NoisySGD(
            noise_scale=0.0,
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size,
            epochs=self.epochs, 
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )

        composed_mechanism = ComposedNoisyMechanism(
            noise_scale=0.0,
            mechanism_list=[
                self.adj_encoder_noisy_sgd,
                self.pma_mechanism,
                self.classifier_noisy_sgd,
            ]
        )

        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_train_nodes)))
                console.info('delta = %.0e' % delta)
            
            self.noise_scale = composed_mechanism.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')
        
        self._adj_encoder = self.adj_encoder_noisy_sgd.prepare_module(self._adj_encoder)
        self._classifier = self.classifier_noisy_sgd.prepare_module(self._classifier)

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '', p: float = 0.6) -> Metrics:
        num_train_nodes = data.train_mask.sum().item()

        if num_train_nodes != self.num_train_nodes:
            self.num_train_nodes = num_train_nodes
            self.calibrate() 
        
        # bound degree here.
        with console.status('bounding the number of neighbors per node'):
            data.adj_t = data.adj_t.t()
            data = BoundOutDegree(self.max_degree)(data)
            data.adj_t = data.adj_t.t()
            assert torch.all(data.adj_t.t().sum(0) <= self.max_degree) # ensure each column of A has at most max_degree nnz.
        return super().fit(data, learn=learn, prefix=prefix)

    def perturb(self, x) -> torch.Tensor:
        # Note that the node to be replaced will be take care of by the following DPSGD. 
        # Hence, the effective sensitivity is sqrt(2*max_degree)
        x = self.pma_mechanism(x, sensitivity=self.norm_scale*np.sqrt(2*self.max_degree))
        return x

    def configure_adj_encoder_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_adj_encoder_optimizer()
        # noise_factor = 2 for replacement dp, replacing 1 node resulting at most max_degree+1 row changes.
        optimizer = self.adj_encoder_noisy_sgd.prepare_optimizer(optimizer, 2.0*np.sqrt(self.max_degree+1))
        return optimizer
    
    def configure_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_optimizer()
        optimizer = self.classifier_noisy_sgd.prepare_optimizer(optimizer, 2.0) # replacement DP
        return optimizer
    
    def adj_data_loader(self,data: Data, stage: Stage) -> SparseAdjRowLoader:
        dataloader = super().adj_data_loader(data,stage)
        if stage == 'train':
            dataloader.wor_sampling = True
            # dataloader.poisson_sampling = True
        return dataloader
    
    def data_loader(self, data: Data, stage: Stage) -> NodeDataLoader:
        dataloader = super().data_loader(data, stage)
        if stage == 'train':
            dataloader.wor_sampling = True
            # dataloader.poisson_sampling = True
        return dataloader

core/methods/node/dpdgc/dpdgc_ndp.py
thupchnsky/dp-gnn
import numpy as np
import torch
from typing import Annotated, Literal, Union
from torch_geometric.data import Data
from torch_sparse import SparseTensor, matmul
from core import console
from core.args.utils import ArgInfo
from core.methods.node import GAP
from core.privacy.algorithms import PMA
from core.modules.base import Metrics


class EdgePrivGAP (GAP):
    """edge-private GAP method"""

    def __init__(self,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[GAP])]
                 ):

        super().__init__(num_classes, **kwargs)
        self.epsilon = epsilon
        self.delta = delta
        self.num_edges = None  # will be used to set delta if it is 'auto'

    def calibrate(self):
        self.pma_mechanism = PMA(noise_scale=0.0, hops=self.hops)
        
        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_edges)))
                console.info('delta = %.0e' % delta)
            
            self.noise_scale = self.pma_mechanism.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')


    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '') -> Metrics:
        if data.num_edges != self.num_edges:
            self.num_edges = data.num_edges
            self.calibrate()

        return super().fit(data, learn=learn, prefix=prefix)

    def _aggregate(self, x: torch.Tensor, adj_t: SparseTensor) -> torch.Tensor:
        x = matmul(adj_t, x)
        x = self.pma_mechanism(x, sensitivity=1)
        return x
core/methods/node/gap/gap_edp.py
thupchnsky/dp-gnn
import torch
from typing import Annotated, Optional
import torch.nn.functional as F
from torch.optim import Adam, SGD, Optimizer
from torch_geometric.data import Data
from torch_geometric.utils import subgraph
from torch_geometric.transforms import ToSparseTensor
from torch_sparse import SparseTensor, matmul
from core import console
from core.args.utils import ArgInfo
from core.methods.node.base import NodeClassification
from core.models.multi_mlp import MultiMLP
from core.modules.base import Metrics
from core.modules.node.cm import ClassificationModule
from core.modules.node.em import EncoderModule


class GAP (NodeClassification):
    """Non-private GAP method"""

    supported_activations = {
        'relu': torch.relu_,
        'selu': torch.selu_,
        'tanh': torch.tanh,
    }

    def __init__(self,
                 num_classes,
                 hops:            Annotated[int,   ArgInfo(help='number of hops', option='-k')] = 2,
                 hidden_dim:      Annotated[int,   ArgInfo(help='dimension of the hidden layers')] = 16,
                 encoder_layers:  Annotated[int,   ArgInfo(help='number of encoder MLP layers')] = 2,
                 base_layers:     Annotated[int,   ArgInfo(help='number of base MLP layers')] = 1,
                 head_layers:     Annotated[int,   ArgInfo(help='number of head MLP layers')] = 1,
                 combine:         Annotated[str,   ArgInfo(help='combination type of transformed hops', choices=MultiMLP.supported_combinations)] = 'cat',
                 activation:      Annotated[str,   ArgInfo(help='type of activation function', choices=supported_activations)] = 'selu',
                 dropout:         Annotated[float, ArgInfo(help='dropout rate')] = 0.0,
                 batch_norm:      Annotated[bool,  ArgInfo(help='if true, then model uses batch normalization')] = True,
                 encoder_epochs:  Annotated[int,   ArgInfo(help='number of epochs for encoder pre-training (ignored if encoder_layers=0)')] = 100,
                 **kwargs:        Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[NodeClassification])]
                 ):

        super().__init__(num_classes, **kwargs)

        if encoder_layers == 0 and encoder_epochs > 0:
            console.warning('encoder_layers is 0, setting encoder_epochs to 0')
            encoder_epochs = 0

        self.hops = hops
        self.encoder_layers = encoder_layers
        self.encoder_epochs = encoder_epochs
        activation_fn = self.supported_activations[activation]
        
        
        self._encoder = EncoderModule(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            encoder_layers=encoder_layers,
            head_layers=1,
            normalize=True,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
        )

        self._classifier = ClassificationModule(
            num_channels=hops+1,
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            base_layers=base_layers,
            head_layers=head_layers,
            combination=combine,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
        )

    @property
    def classifier(self) -> ClassificationModule:
        return self._classifier

    def reset_parameters(self):
        self._encoder.reset_parameters()
        super().reset_parameters()

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '') -> Metrics:
        self.learning_setting = learn
        
        if learn == 'transductive':
            
            self.data = data.to(self.device, non_blocking=True)

            # pre-train encoder
            if self.encoder_layers > 0:
                self.data = self.pretrain_encoder(self.data, prefix=prefix)

            # compute aggregations
            self.data = self.compute_aggregations(self.data)

            # train classifier
            return super().fit(self.data, learn=learn, prefix=prefix)
        
        elif learn == 'inductive':
            
            # we first record the clone of data in self.data
            self.data = data.clone().detach().to(self.device, non_blocking=True)
            
            # in inductive setting we only know the information about the nodes in training set
            # the input data contains the complete graph information so we need to first extract the training dataset
            data = data.to(self.device, non_blocking=True)
            
            # extract the training graph (train_mask + val_mask)
            node_indices = torch.LongTensor([i for i in range(data.x.size(0))]).to(self.device)
            train_node_indices = node_indices[data.train_mask | data.val_mask]
            row, col, _ = data.adj_t.t().coo()
            edge_index = torch.stack([row, col], dim=0)
            
            train_edge_index, _ = subgraph(train_node_indices, edge_index, relabel_nodes=False)
            data.edge_index = train_edge_index  # change the edge index
            data = ToSparseTensor()(data)  # now data.adj_t has changed
            
            # pre-train encoder
            if self.encoder_layers > 0:
                data = self.pretrain_encoder(data, prefix=prefix)

            # compute aggregations
            data = self.compute_aggregations(data)

            # train classifier
            return super().fit(data, learn=learn, prefix=prefix)
        
        else:
            
            raise Exception('Not supported learning setting.')
    
    def test(self, data: Optional[Data] = None, prefix: str = '') -> Metrics:
        if (data is None or data == self.data) and self.learning_setting == 'transductive':
            data = self.data
        else:
            data = data.to(self.device, non_blocking=True)
            data.x = self._encoder.predict(data)
            data = self.compute_aggregations(data)

        return super().test(data, prefix=prefix)

    def predict(self, data: Optional[Data] = None) -> torch.Tensor:
        if (data is None or data == self.data) and self.learning_setting == 'transductive':
            data = self.data
        else:
            data.x = self._encoder.predict(data)
            data = self.compute_aggregations(data)

        return super().predict(data)

    def _aggregate(self, x: torch.Tensor, adj_t: SparseTensor) -> torch.Tensor:
        return matmul(adj_t, x)

    def _normalize(self, x: torch.Tensor) -> torch.Tensor:
        return F.normalize(x, p=2, dim=-1)

    def pretrain_encoder(self, data: Data, prefix: str) -> Data:
        console.info('pretraining encoder')
        self._encoder.to(self.device)
        
        self.trainer.fit(
            model=self._encoder,
            epochs=self.encoder_epochs,
            optimizer=self.configure_encoder_optimizer(), 
            train_dataloader=self.data_loader(data, 'train'), 
            val_dataloader=self.data_loader(data, 'val'),
            test_dataloader=None,
            checkpoint=True,
            prefix=f'{prefix}encoder/',
        )

        self.trainer.reset()
        data.x = self._encoder.predict(data)
        return data

    def compute_aggregations(self, data: Data) -> Data:
        with console.status('computing aggregations'):
            x = F.normalize(data.x, p=2, dim=-1)
            x_list = [x]

            for _ in range(self.hops):
                x = self._aggregate(x, data.adj_t)
                x = self._normalize(x)
                x_list.append(x)

            data.x = torch.stack(x_list, dim=-1)
        return data

    def configure_encoder_optimizer(self) -> Optimizer:
        Optim = {'sgd': SGD, 'adam': Adam}[self.optimizer_name]
        return Optim(self._encoder.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
core/methods/node/gap/gap_inf.py
thupchnsky/dp-gnn
import numpy as np
import torch
from typing import Annotated, Literal, Union
from torch_geometric.data import Data
from torch_sparse import SparseTensor, matmul
from opacus.optimizers import DPOptimizer
from core import console
from core.args.utils import ArgInfo
from core.data.loader import NodeDataLoader
from core.methods.node import GAP
from core.privacy.mechanisms import ComposedNoisyMechanism
from core.privacy.algorithms import PMA, NoisySGD
from core.data.transforms import BoundOutDegree
from core.modules.base import Metrics, Stage


class NodePrivGAP (GAP):
    """node-private GAP method"""

    def __init__(self,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 max_degree:    Annotated[int,   ArgInfo(help='max degree to sample per each node')] = 100,
                 max_grad_norm: Annotated[float, ArgInfo(help='maximum norm of the per-sample gradients')] = 1.0,
                 batch_size:    Annotated[int,   ArgInfo(help='batch size')] = 256,
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[GAP], exclude=['batch_norm'])]
                 ):

        super().__init__(num_classes, 
            batch_norm=False, 
            batch_size=batch_size, 
            **kwargs
        )
        self.epsilon = epsilon
        self.delta = delta
        self.max_degree = max_degree
        self.max_grad_norm = max_grad_norm

        self.num_train_nodes = None  # will be used to set delta if it is 'auto'

    def calibrate(self):
        self.pma_mechanism = PMA(noise_scale=0.0, hops=self.hops)

        self.encoder_noisy_sgd = NoisySGD(
            noise_scale=0.0, 
            dataset_size=self.num_train_nodes, 
            batch_size=self.batch_size, 
            epochs=self.encoder_epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )

        self.classifier_noisy_sgd = NoisySGD(
            noise_scale=0.0, 
            dataset_size=self.num_train_nodes, 
            batch_size=self.batch_size, 
            epochs=self.epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )

        composed_mechanism = ComposedNoisyMechanism(
            noise_scale=0.0,
            mechanism_list=[
                self.encoder_noisy_sgd, 
                self.pma_mechanism, 
                self.classifier_noisy_sgd
            ]
        )

        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_train_nodes)))
                console.info('delta = %.0e' % delta)
            
            self.noise_scale = composed_mechanism.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')

        self._encoder = self.encoder_noisy_sgd.prepare_module(self._encoder)
        self._classifier = self.classifier_noisy_sgd.prepare_module(self._classifier)

    def fit(self, data: Data, learn: str = 'transductive', prefix: str = '') -> Metrics:
        num_train_nodes = data.train_mask.sum().item()

        if num_train_nodes != self.num_train_nodes:
            self.num_train_nodes = num_train_nodes
            self.calibrate()

        return super().fit(data, learn=learn, prefix=prefix)

    def compute_aggregations(self, data: Data) -> Data:
        with console.status('bounding the number of neighbors per node'):
            data = BoundOutDegree(self.max_degree)(data)
            assert torch.all(data.adj_t.sum(0)<=self.max_degree)
        return super().compute_aggregations(data)

    def _aggregate(self, x: torch.Tensor, adj_t: SparseTensor) -> torch.Tensor:
        x = matmul(adj_t, x)
        # Note that the node to be replaced will be treated seperately. Hence, 
        # the effective sensitivity is 2*sqrt(max_degree)
        x = self.pma_mechanism(x, sensitivity=2*np.sqrt(self.max_degree))
        return x

    def data_loader(self, data: Data, stage: Stage) -> NodeDataLoader:
        dataloader = super().data_loader(data, stage)
        if stage == 'train':
            dataloader.wor_sampling = True
        return dataloader

    def configure_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_optimizer()
        optimizer = self.classifier_noisy_sgd.prepare_optimizer(optimizer, 2.0)
        return optimizer

    def configure_encoder_optimizer(self) -> DPOptimizer:
        optimizer = super().configure_encoder_optimizer()
        optimizer = self.encoder_noisy_sgd.prepare_optimizer(optimizer, 2.0)
        return optimizer

core/methods/node/gap/gap_ndp.py
thupchnsky/dp-gnn
import torch
from typing import Annotated
from core.args.utils import ArgInfo
from core.methods.node import NodeClassification
from core.modules.node.mlp import MLPNodeClassifier


class MLP (NodeClassification):
    """Non-private MLP method"""

    supported_activations = {
        'relu': torch.relu_,
        'selu': torch.selu_,
        'tanh': torch.tanh,
    }

    def __init__(self,
                 num_classes,
                 hidden_dim:      Annotated[int,   ArgInfo(help='dimension of the hidden layers')] = 16,
                 num_layers:      Annotated[int,   ArgInfo(help='number of MLP layers')] = 2,
                 activation:      Annotated[str,   ArgInfo(help='type of activation function', choices=supported_activations)] = 'selu',
                 dropout:         Annotated[float, ArgInfo(help='dropout rate')] = 0.0,
                 batch_norm:      Annotated[bool,  ArgInfo(help='if true, then model uses batch normalization')] = True,
                 **kwargs:        Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[NodeClassification])]
                 ):

        assert num_layers > 1, 'number of layers must be greater than 1'
        super().__init__(num_classes, **kwargs)

        self.num_layers = num_layers
        activation_fn = self.supported_activations[activation]

        self._classifier = MLPNodeClassifier(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
        )

    @property
    def classifier(self) -> MLPNodeClassifier:
        return self._classifier

core/methods/node/mlp/mlp.py
thupchnsky/dp-gnn
import numpy as np
from typing import Annotated, Literal, Union
from torch.optim import Optimizer
from torch_geometric.data import Data
from core import console
from core.args.utils import ArgInfo
from core.methods.node import MLP
from core.privacy.algorithms.noisy_sgd import NoisySGD
from core.modules.base import Metrics, Stage
from core.data.loader import NodeDataLoader


class PrivMLP (MLP):
    """node-private MLP method"""

    def __init__(self,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 max_grad_norm: Annotated[float, ArgInfo(help='maximum norm of the per-sample gradients')] = 1.0,
                 batch_size:    Annotated[int,   ArgInfo(help='batch size')] = 256,
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[MLP], exclude=['batch_norm'])]
                 ):

        super().__init__(num_classes, 
            batch_norm=False, 
            batch_size=batch_size, 
            **kwargs
        )

        self.epsilon = epsilon
        self.delta = delta
        self.max_grad_norm = max_grad_norm
        self.num_train_nodes = None         # will be used to auto set delta
        
    def calibrate(self):
        self.noisy_sgd = NoisySGD(
            noise_scale=0.0, 
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size, 
            epochs=self.epochs,
            max_grad_norm=self.max_grad_norm,
            replacement=True
        )

        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_train_nodes)))
                console.info('delta = %.0e', delta)
            
            self.noise_scale = self.noisy_sgd.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')

        self._classifier = self.noisy_sgd.prepare_module(self._classifier)

    def fit(self, data: Data, prefix: str = '') -> Metrics:
        num_train_nodes = data.train_mask.sum().item()

        if num_train_nodes != self.num_train_nodes:
            self.num_train_nodes = num_train_nodes
            self.calibrate()

        return super().fit(data, prefix=prefix)

    def data_loader(self, data: Data, stage: Stage) -> NodeDataLoader:
        dataloader = super().data_loader(data, stage)
        if stage == 'train':
            dataloader.wor_sampling = True
        return dataloader

    def configure_optimizer(self) -> Optimizer:
        optimizer = super().configure_optimizer()
        optimizer = self.noisy_sgd.prepare_optimizer(optimizer, 2.0) # replacement DP
        return optimizer

core/methods/node/mlp/mlp_dp.py
thupchnsky/dp-gnn
from typing import Annotated, Optional
import torch
from torch_geometric.data import Data
from core import console
from core.args.utils import ArgInfo
from core.methods.node import SAGE
from core.privacy.algorithms import AsymmetricRandResponse
from core.modules.base import Metrics


class EdgePrivSAGE (SAGE):
    """edge-private GraphSAGE method"""
    
    def __init__(self,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[SAGE])]
                 ):

        super().__init__(num_classes, **kwargs)
        self.mechanism = AsymmetricRandResponse(eps=epsilon)

    def perturb_data(self, data: Data) -> Data:
        data.adj_t = data.adj_t.to(self.device)
        with console.status('perturbing graph structure'):
            data.adj_t = self.mechanism(data.adj_t, chunk_size=500)
        return data

    def fit(self, data: Data, prefix: str = '') -> Metrics:
        data = self.perturb_data(data)
        return super().fit(data, prefix=prefix)

    def test(self, data: Optional[Data] = None, prefix: str = '') -> Metrics:
        if data is not None and data != self.data:
            data = self.perturb_data(data)
        return super().test(data, prefix=prefix)

    def predict(self, data: Optional[Data] = None) -> torch.Tensor:
        if data is not None and data != self.data:
            data = self.perturb_data(data)
        return super().predict(data)

core/methods/node/sage/sage_edp.py
thupchnsky/dp-gnn
import torch
from typing import Annotated
from core.args.utils import ArgInfo
from core.methods.node import NodeClassification
from core.modules.node.sage import SAGENodeClassifier


class SAGE (NodeClassification):
    """Non-private GraphSAGE method"""
    
    supported_activations = {
        'relu': torch.relu_,
        'selu': torch.selu_,
        'tanh': torch.tanh,
    }

    def __init__(self,
                 num_classes,
                 hidden_dim:      Annotated[int,   ArgInfo(help='dimension of the hidden layers')] = 16,
                 base_layers:     Annotated[int,   ArgInfo(help='number of base MLP layers')] = 2,
                 mp_layers:       Annotated[int,   ArgInfo(help='number of GNN layers')] = 1,
                 head_layers:     Annotated[int,   ArgInfo(help='number of head MLP layers')] = 1,
                 activation:      Annotated[str,   ArgInfo(help='type of activation function', choices=supported_activations)] = 'selu',
                 dropout:         Annotated[float, ArgInfo(help='dropout rate')] = 0.0,
                 batch_norm:      Annotated[bool,  ArgInfo(help='if true, then model uses batch normalization')] = True,
                 **kwargs:        Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[NodeClassification])]
                 ):

        assert mp_layers >= 1, 'number of message-passing layers must be at least 1'
        super().__init__(num_classes=num_classes, **kwargs)
        
        self.base_layers = base_layers
        self.mp_layers = mp_layers
        self.head_layers = head_layers
        activation_fn = self.supported_activations[activation]

        self._classifier = SAGENodeClassifier(
            hidden_dim=hidden_dim, 
            num_classes=num_classes, 
            base_layers=base_layers,
            mp_layers=mp_layers, 
            head_layers=head_layers, 
            normalize=False,
            activation_fn=activation_fn, 
            dropout=dropout, 
            batch_norm=batch_norm,
        )

    @property
    def classifier(self) -> SAGENodeClassifier:
        return self._classifier

core/methods/node/sage/sage_inf.py
thupchnsky/dp-gnn
import numpy as np
from typing import Annotated, Optional, Union, Literal
import torch
from torch.optim import Optimizer
from torch_geometric.data import Data
from core import console
from core.args.utils import ArgInfo
from core.data.loader import NodeDataLoader
from core.data.transforms import BoundDegree
from core.methods.node import SAGE
from core.modules.base import Metrics, Stage
from core.privacy.algorithms import GNNBasedNoisySGD
from core.privacy.mechanisms import GaussianMechanism
from core.privacy.mechanisms import ComposedNoisyMechanism


class NodePrivSAGE (SAGE):
    """node-private GraphSAGE method"""
    
    def __init__(self,
                 num_classes,
                 epsilon:       Annotated[float, ArgInfo(help='DP epsilon parameter', option='-e')],
                 delta:         Annotated[Union[Literal['auto'], float], 
                                                 ArgInfo(help='DP delta parameter (if "auto", sets a proper value based on data size)', option='-d')] = 'auto',
                 max_degree:    Annotated[int,   ArgInfo(help='max degree to sample per each node')] = 100,
                 max_grad_norm: Annotated[float, ArgInfo(help='maximum norm of the per-sample gradients')] = 1.0,
                 batch_size:    Annotated[int,   ArgInfo(help='batch size')] = 256,
                 output_gdp:    Annotated[int,   ArgInfo(help='if 1 (true), then model add noise at output to ensure GDP')] = 1,
                 **kwargs:      Annotated[dict,  ArgInfo(help='extra options passed to base class', bases=[SAGE], exclude=['batch_norm', 'mp_layers', 'val_interval'])]
                 ):

        super().__init__(
            num_classes=num_classes, 
            batch_size=batch_size, 
            batch_norm=False, 
            mp_layers=1, 
            val_interval=0,
            **kwargs
        )
        self.epsilon = epsilon
        self.delta = delta
        self.max_degree = max_degree
        self.max_grad_norm = max_grad_norm
        
        self.num_train_nodes = None         # will be used to auto set delta
        self.classifier.normalize = True    # required to bound sensitivity
        
        if output_gdp == 1:
            self.output_gdp = True
        else:
            self.output_gdp = False
        if not self.output_gdp:
            self.classifier.normalize = False
            
    def calibrate(self):
        self.noisy_sgd = GNNBasedNoisySGD(
            noise_scale=0.0, 
            dataset_size=self.num_train_nodes,
            batch_size=self.batch_size, 
            epochs=self.epochs,
            max_grad_norm=self.max_grad_norm,
            max_degree=self.max_degree,
        )
        
        
        self.noisy_aggr_gm = GaussianMechanism(noise_scale=0.0)
        
        if self.output_gdp:
            composed_mechanism = ComposedNoisyMechanism(
                noise_scale=0.0,
                mechanism_list=[self.noisy_sgd, self.noisy_aggr_gm], 
                coeff_list=[1,1]
            )
        else:
            composed_mechanism = ComposedNoisyMechanism(
                noise_scale=0.0,
                mechanism_list=[self.noisy_sgd], 
            )

        if hasattr(self, 'noisy_aggr_hook'):
            self.noisy_aggr_hook.remove()
        
        
        if self.output_gdp:
            self.noisy_aggr_hook = self.classifier.gnn.convs[0].register_message_and_aggregate_forward_hook(
                lambda module, inputs, output: 
                    self.noisy_aggr_gm(data=output, sensitivity=2*np.sqrt(self.max_degree)) if not module.training else output # for replacement, 2*np.sqrt(self.max_degree).
            )
        else:
            self.noisy_aggr_hook = self.classifier.gnn.convs[0].register_message_and_aggregate_forward_hook(
                lambda module, inputs, output: 
                    output if not module.training else output # Do not add output perturbation.
            )

        with console.status('calibrating noise to privacy budget'):
            if self.delta == 'auto':
                delta = 0.0 if np.isinf(self.epsilon) else 1. / (10 ** len(str(self.num_train_nodes)))
                console.info('delta = %.0e' % delta)
            
            self.noise_scale = composed_mechanism.calibrate(eps=self.epsilon, delta=delta)
            console.info(f'noise scale: {self.noise_scale:.4f}\n')

        self._classifier = self.noisy_sgd.prepare_module(self._classifier)

    def sample_neighbors(self, data: Data) -> Data:
        data = data.to(self.device, non_blocking=True)
        with console.status('bounding the number of neighbors per node'):
            data = BoundDegree(self.max_degree)(data)
        return data

    def fit(self, data: Data, prefix: str = '') -> Metrics:
        num_train_nodes = data.train_mask.sum().item()

        if num_train_nodes != self.num_train_nodes:
            self.num_train_nodes = num_train_nodes
            self.calibrate()

        data = self.sample_neighbors(data)
        return super().fit(data, prefix=prefix)

    def test(self, data: Optional[Data] = None, prefix: str = '') -> Metrics:
        if data is not None and data != self.data:
            data = self.sample_neighbors(data)
        return super().test(data, prefix=prefix)

    def predict(self, data: Optional[Data] = None) -> torch.Tensor:
        if data is not None and data != self.data:
            data = self.sample_neighbors(data)
        return super().predict(data)

    def data_loader(self, data: Data, stage: Stage) -> NodeDataLoader:
        dataloader = super().data_loader(data, stage)
        if stage == 'train':
            dataloader.hops = 1
            dataloader.poisson_sampling = True
        return dataloader

    def configure_optimizer(self) -> Optimizer:
        optimizer = super().configure_optimizer()
        optimizer = self.noisy_sgd.prepare_optimizer(optimizer)
        return optimizer

core/methods/node/sage/sage_ndp.py
thupchnsky/dp-gnn
from core.models.mlp import MLP

core/models/__init__.py
thupchnsky/dp-gnn
from typing import Callable, Iterable, Literal, get_args
import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import ModuleList, LazyBatchNorm1d, BatchNorm1d, Linear, Module, Dropout
from core.models import MLP


class LINKX(Module):
    CombType = Literal['cat', 'sum', 'max', 'mean', 'att']
    supported_combinations = get_args(CombType)

    def __init__(self, *,
                 num_features: int,
                 output_dim: int,
                 hidden_dim: int = 16,
                 base_layers: int = 2,
                 head_layers: int = 1,
                 combination: CombType = 'cat',
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 dropout: float = 0.0, 
                 batch_norm: bool = False,
                 plain_last: bool = True,
                 ):

        super().__init__()
        
        self.num_features = num_features
        self.combination = combination
        self.activation_fn = activation_fn
        self.dropout_fn = Dropout(dropout, inplace=True)

        self.mlp_feat = MLP(hidden_dim=hidden_dim,
                            output_dim=hidden_dim,
                            num_layers=base_layers,
                            dropout=dropout,
                            activation_fn=activation_fn,
                            batch_norm=batch_norm,
                            plain_last=True)
        
        self.mlp_W = MLP(hidden_dim=hidden_dim,
                         output_dim=hidden_dim,
                         num_layers=base_layers,
                         dropout=dropout,
                         activation_fn=activation_fn,
                         batch_norm=batch_norm,
                         plain_last=True)
        
        if combination == 'att':
            self.hidden_dim = hidden_dim
            self.num_heads = 2
            self.Q = Linear(in_features=hidden_dim, out_features=self.num_heads, bias=False)

        self.bn1 = BatchNorm1d(2*hidden_dim) if batch_norm else False
        self.bn2 = BatchNorm1d(hidden_dim) if batch_norm else False

        self.head_mlp = MLP(
            output_dim=output_dim,
            hidden_dim=hidden_dim,
            num_layers=head_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
            plain_last=plain_last,
        )

    def forward(self, x: Tensor) -> Tensor:  
        # x_stack input is of form [raw features, embedded adj]
        xX = self.mlp_feat(x[:, 0: self.num_features])
        xA = x[:, self.num_features:]
        
        h_list = [xX, xA]
        h = self.combine(h_list)
        
        h = self.bn1(h) if self.bn1 else h
        h = self.dropout_fn(h)
        h = self.activation_fn(h)
        
        h = self.mlp_W(h)
        h = h + xX + xA
        
        h = self.bn2(h) if self.bn2 else h
        h = self.dropout_fn(h)
        h = self.activation_fn(h)
        
        h = self.head_mlp(h)
        return h

    def combine(self, h_list: Iterable[Tensor]) -> Tensor:
        if self.combination == 'cat':
            return torch.cat(h_list, dim=-1)
        elif self.combination == 'sum':
            return torch.stack(h_list, dim=0).sum(dim=0)
        elif self.combination == 'mean':
            return torch.stack(h_list, dim=0).mean(dim=0)
        elif self.combination == 'max':
            return torch.stack(h_list, dim=0).max(dim=0).values
        elif self.combination == 'att':
            H = torch.stack(h_list, dim=1)  # (node, hop, dim)
            W = F.leaky_relu(self.Q(H), 0.2).softmax(dim=0)  # (node, hop, head)
            out = H.transpose(1, 2).matmul(W).view(-1, self.hidden_dim * self.num_heads)
            return out
        else:
            raise ValueError(f'Unknown combination type {self.combination}')

    def reset_parameters(self):
        self.mlp_feat.reset_parameters()
        self.mlp_W.reset_parameters()
        if self.combination == 'att':
            self.Q.reset_parameters()
        if self.bn1:
            self.bn1.reset_parameters()
        if self.bn2:
            self.bn2.reset_parameters()
        self.head_mlp.reset_parameters()

core/models/linkx.py
thupchnsky/dp-gnn
from typing import Callable
import torch
from torch import Tensor
from torch.nn import Dropout, ModuleList, BatchNorm1d, Module
from torch_geometric.nn import Linear
from opacus.grad_sample import register_grad_sampler, compute_linear_grad_sample


@register_grad_sampler(Linear)
def compute_lazy_linear_grad_sample(layer, activations, backprops):
    return compute_linear_grad_sample(layer, activations, backprops)


class MLP(Module):
    """
    A multi-layer perceptron (MLP) model.
    This implementation handles 0-layer configurations as well.
    """
    def __init__(self, *,
                 output_dim: int,
                 hidden_dim: int = 16,  
                 num_layers: int = 2, 
                 dropout: float = 0.0, 
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 plain_last: bool = True,
                 ):
        super().__init__()
        self.num_layers = num_layers
        self.dropout_fn = Dropout(dropout, inplace=True)
        self.activation_fn = activation_fn
        self.plain_last = plain_last

        dimensions = [hidden_dim] * (num_layers - 1) + [output_dim] * (num_layers > 0)
        self.layers: list[Linear] = ModuleList([Linear(-1, dim) for dim in dimensions])
        
        num_bns = batch_norm * (num_layers - int(plain_last))
        self.bns: list[BatchNorm1d] = []
        if batch_norm:
            self.bns = ModuleList([BatchNorm1d(hidden_dim) for _ in range(num_bns)])
        
    def forward(self, x: Tensor) -> Tensor:
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < self.num_layers - self.plain_last:
                x = self.bns[i](x) if self.bns else x
                x = self.dropout_fn(x)
                x = self.activation_fn(x)
        return x

    def reset_parameters(self):
        for layer in self.layers:
            layer.reset_parameters()
        
        for bn in self.bns:
            bn.reset_parameters()
core/models/mlp.py
thupchnsky/dp-gnn
from typing import Callable, Iterable, Literal, get_args
import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import ModuleList, LazyBatchNorm1d, Linear, Module, Dropout
from core.models import MLP


class MultiMLP(Module):
    CombType = Literal['cat', 'sum', 'max', 'mean', 'att']
    supported_combinations = get_args(CombType)

    def __init__(self, *, 
                 num_channels: int, 
                 output_dim: int,
                 hidden_dim: int = 16,  
                 base_layers: int = 2, 
                 head_layers: int = 1, 
                 combination: CombType = 'cat',
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 dropout: float = 0.0, 
                 batch_norm: bool = False,
                 plain_last: bool = True,
                 ):

        super().__init__()

        self.combination = combination
        self.activation_fn = activation_fn
        self.dropout_fn = Dropout(dropout, inplace=True)

        self.base_mlps: list[MLP] = ModuleList([
            MLP(
                hidden_dim=hidden_dim,
                output_dim=hidden_dim,
                num_layers=base_layers,
                dropout=dropout,
                activation_fn=activation_fn,
                batch_norm=batch_norm,
                plain_last=True,
            ) for _ in range(num_channels)]
        )

        if combination == 'att':
            self.hidden_dim = hidden_dim
            self.num_heads = num_channels
            self.Q = Linear(in_features=hidden_dim, out_features=self.num_heads, bias=False)

        self.bn = LazyBatchNorm1d() if batch_norm else False

        self.head_mlp = MLP(
            output_dim=output_dim,
            hidden_dim=hidden_dim,
            num_layers=head_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
            plain_last=plain_last,
        )

    def forward(self, x_stack: Tensor) -> Tensor:
        x_stack = x_stack.permute(2, 0, 1) # (hop, node, input_dim)
        h_list = [mlp(x) for x, mlp in zip(x_stack, self.base_mlps)]
        h = self.combine(h_list)
        h = F.normalize(h, p=2, dim=-1)
        h = self.bn(h) if self.bn else h
        h = self.dropout_fn(h)
        h = self.activation_fn(h)
        h = self.head_mlp(h)
        return h

    def combine(self, h_list: Iterable[Tensor]) -> Tensor:
        if self.combination == 'cat':
            return torch.cat(h_list, dim=-1)
        elif self.combination == 'sum':
            return torch.stack(h_list, dim=0).sum(dim=0)
        elif self.combination == 'mean':
            return torch.stack(h_list, dim=0).mean(dim=0)
        elif self.combination == 'max':
            return torch.stack(h_list, dim=0).max(dim=0).values
        elif self.combination == 'att':
            H = torch.stack(h_list, dim=1)  # (node, hop, dim)
            W = F.leaky_relu(self.Q(H), 0.2).softmax(dim=0)  # (node, hop, head)
            out = H.transpose(1, 2).matmul(W).view(-1, self.hidden_dim * self.num_heads)
            return out
        else:
            raise ValueError(f'Unknown combination type {self.combination}')

    def reset_parameters(self):
        for mlp in self.base_mlps: mlp.reset_parameters()
        if self.combination == 'att':
            self.Q.reset_parameters()
        if self.bn: self.bn.reset_parameters()
        self.head_mlp.reset_parameters()

core/models/multi_mlp.py
thupchnsky/dp-gnn
from typing import Callable
import torch
from torch import Tensor
from torch.nn import Dropout, BatchNorm1d
from torch_geometric.nn import GraphSAGE
from torch_sparse import SparseTensor


class SAGE(GraphSAGE):
    """
    A GraphSAGE model.
    This implementation supports plain_last option.
    """
    def __init__(self, *,
                 output_dim: int,
                 hidden_dim: int = 16,  
                 num_layers: int = 2, 
                 dropout: float = 0.0, 
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 plain_last: bool = True,
                 ):
        super().__init__(
            in_channels=-1,
            hidden_channels=hidden_dim,
            num_layers=num_layers,
            out_channels=output_dim,
            dropout=dropout,
            act=activation_fn,
            norm=BatchNorm1d(hidden_dim) if batch_norm else None,
            jk=None,
            aggr='sum',
            root_weight=True,
            normalize=True,
        )
        
        self.dropout_fn = Dropout(dropout, inplace=True)
        self.activation_fn = activation_fn
        self.batch_norm = batch_norm
        self.plain_last = plain_last
        if not plain_last and batch_norm:
            self.bn = BatchNorm1d(output_dim)
        
    def forward(self, x: Tensor, adj_t: SparseTensor) -> Tensor:
        x = super().forward(x, adj_t)
        if not self.plain_last:
            x = self.bn(x) if self.batch_norm else x
            x = self.dropout_fn(x)
            x = self.activation_fn(x)
        return x        

    def reset_parameters(self):
        super().reset_parameters()
        if not self.plain_last and self.batch_norm:
            self.bn.reset_parameters()
core/models/sage.py
thupchnsky/dp-gnn
from core.modules.base import TrainableModule
core/modules/__init__.py
thupchnsky/dp-gnn
from typing import Literal, Optional
from torch import Tensor
from torch.nn import Module
from torch.types import Number
from torch_geometric.data import Data
from abc import ABC, abstractmethod


Stage = Literal['train', 'val', 'test']
Metrics = dict[str, Number]


class TrainableModule(Module, ABC):
    @abstractmethod
    def forward(self, *args, **kwargs): pass

    @abstractmethod
    def step(self, data: Data, stage: Stage) -> tuple[Optional[Tensor], Metrics]: pass

    @abstractmethod
    def predict(self, data: Data) -> Tensor: pass

    @abstractmethod
    def reset_parameters(self): pass

core/modules/base.py
thupchnsky/dp-gnn
from typing import Callable
import torch
from torch import Tensor
from core.models.multi_mlp import MultiMLP
from core.models.linkx import LINKX
from core.modules.node.mlp import MLPNodeClassifier


class ClassificationModule(MLPNodeClassifier):
    def __init__(self, *,
                 num_channels: int, 
                 num_classes: int,
                 hidden_dim: int = 16,  
                 base_layers: int = 2, 
                 head_layers: int = 1, 
                 combination: MultiMLP.CombType = 'cat',
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 dropout: float = 0.0, 
                 batch_norm: bool = False,
                 ):

        super().__init__(num_classes=num_classes)  # this is dummy, but necessary

        self.model = MultiMLP(
            num_channels=num_channels,
            output_dim=num_classes,
            hidden_dim=hidden_dim,
            base_layers=base_layers,
            head_layers=head_layers,
            combination=combination,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
            plain_last=True,
        )


class DPDGCClassificationModule(MLPNodeClassifier):
    def __init__(self, *,
                 num_features: int,
                 num_classes: int,
                 hidden_dim: int = 16,  
                 base_layers: int = 2,
                 head_layers: int = 1,
                 combination: MultiMLP.CombType = 'cat',
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 dropout: float = 0.0, 
                 batch_norm: bool = False,
                 ):

        super().__init__(num_classes=num_classes)  # this is dummy, but necessary

        self.model = LINKX(
            num_features=num_features,
            output_dim=num_classes,
            hidden_dim=hidden_dim,
            base_layers=base_layers,
            head_layers=head_layers,
            combination=combination,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
            plain_last=True,
        )

core/modules/node/cm.py
thupchnsky/dp-gnn
from typing import Callable
import math
import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import BatchNorm1d, Dropout
from torch.nn.parameter import Parameter
from core.models import MLP
from torch_geometric.data import Data
from torch_geometric.nn import Linear
from core.modules.node.mlp import MLPNodeClassifier, MLPClassifier
from core.modules.node.graph_mlp import GraphNodeClassifier
from torch_sparse import SparseTensor, matmul


class EncoderModule(MLPNodeClassifier):
    def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,  
                 encoder_layers: int = 2, 
                 head_layers: int = 1, 
                 normalize: bool = True,
                 dropout: float = 0.0, 
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 ):

        super().__init__(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            num_layers=head_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
        )

        self.encoder_mlp = MLP(
            hidden_dim=hidden_dim,
            output_dim=hidden_dim,
            num_layers=encoder_layers,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
            plain_last=True,
        )

        self.dropout_fn = Dropout(p=dropout, inplace=True)
        self.activation_fn = activation_fn
        self.normalize = normalize
        self.bn = BatchNorm1d(hidden_dim) if batch_norm else False

    def forward(self, x: Tensor) -> Tensor:
        x = self.encoder_mlp(x)
        if self.normalize:
            x = F.normalize(x, p=2, dim=-1)
        x = self.bn(x) if self.bn else x
        x = self.dropout_fn(x)
        x = self.activation_fn(x)
        x = super().forward(x)
        return x

    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        x = data.x
        x = self.encoder_mlp(x)
        if self.normalize:
            x = F.normalize(x, p=2, dim=-1)
        return x
        
    def reset_parameters(self):
        super().reset_parameters()
        self.encoder_mlp.reset_parameters()
        if self.bn:
            self.bn.reset_parameters()


class GraphEncoderModule(GraphNodeClassifier):
    def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,
                 encoder_layers: int = 2, 
                 head_layers: int = 1,
                 hops: int = 5,
                 dropout: float = 0.0,
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 mlp_inst_norm: bool = False,
                 ):

        super().__init__(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            num_layers=head_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
        )
        
        self.encoder_mlp = MLP(
            hidden_dim=hidden_dim,
            output_dim=hidden_dim,
            num_layers=encoder_layers,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
            plain_last=True,
        )
        
        self.hops = hops
        self.mlp_inst_norm = mlp_inst_norm
        self.dropout_fn = Dropout(p=dropout, inplace=True)
        self.activation_fn = activation_fn
        self.bn = BatchNorm1d(hidden_dim * (hops + 1)) if batch_norm else False  # [noisy_x_emb, K propagated features]
    
    def forward(self, x: Tensor, adj_t: SparseTensor, noisy_x_emb: Tensor) -> Tensor:
        x = self.encoder_mlp(x)
        if self.mlp_inst_norm:
            x = F.normalize(x, p=2, dim=-1)
        
        Z = [noisy_x_emb]
        for _ in range(self.hops):
            x = matmul(adj_t, x)
            x = F.normalize(x, p=2, dim=-1)
            Z.append(x)
        
        x = torch.cat(Z, dim=-1) 
        x = self.bn(x) if self.bn else x
        x = self.dropout_fn(x)
        x = self.activation_fn(x)
        x = super().forward(x, adj_t, noisy_x_emb)
        return x

    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        x = self.encoder_mlp(data.x)  # we do not perform normalization here
        return x
        
    def reset_parameters(self):
        super().reset_parameters()
        self.encoder_mlp.reset_parameters()
        if self.bn:
            self.bn.reset_parameters()


class AdjEncoderModule(MLPClassifier):
    '''
        This encoder module is used for MLP_A, and it is similar as EncoderModule, except that now 
        we need to normalize the rows of the weight matrix W.
    '''
    def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,
                 head_layers: int = 1,
                 dropout: float = 0.0,
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 batch_norm: bool = False,
                 norm_scale: float = 1.0,
                 ):

        super().__init__(
            num_classes=num_classes,
            hidden_dim=hidden_dim,
            num_layers=head_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
        )
        
        # The MLP_A is hard-coded to be a one layer MLP (XW^T+b)
        self.mlp_adj = Linear(-1, hidden_dim)
        self.norm_scale = norm_scale
        
        # Freeze the head MLP as initialization, and disable bias
        for lin in self.model.layers:
            lin.bias = torch.nn.init.constant_(lin.bias, 0.0)
            # we use random orthonormal matrix to construct projection. It prevents potential rank collapse.
            Rot = torch.randn(hidden_dim,hidden_dim)
            svd = torch.linalg.svd(Rot)
            orth = svd[0] @ svd[2]
            lin.weight = Parameter(torch.eye(num_classes,hidden_dim) @ orth * math.sqrt(3.0/hidden_dim*2))
        for param in self.model.parameters():
            param.requires_grad = False
        
        # extract the weight matrix and do column normalization of W (row normalization of W^T)
        # self.mlp_adj.state_dict()['weight'] = F.normalize(self.mlp_adj.state_dict()['weight'], p=2, dim=0)
        # debug use only
        # print(torch.norm(self.mlp_adj.state_dict()['weight'], dim=0))
        
        self.dropout_fn = Dropout(p=dropout, inplace=True)
        self.activation_fn = activation_fn
        self.bn = BatchNorm1d(hidden_dim) if batch_norm else False

    def forward(self, x: Tensor) -> Tensor:
        # Here the x would actually be data.adj_t
        x = self.mlp_adj(x)
        self.mlp_adj.state_dict()['weight'] = F.normalize(self.mlp_adj.state_dict()['weight'], p=2, dim=0) * self.norm_scale
        x = self.bn(x) if self.bn else x
        x = self.dropout_fn(x)
        x = self.activation_fn(x)
        x = super().forward(x)
        return x
    
    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        # Ensure weight matrix is normalized.
        self.mlp_adj.state_dict()['weight'] = F.normalize(self.mlp_adj.state_dict()['weight'], p=2, dim=0) * self.norm_scale
        x = data.adj_t.t().to_torch_sparse_coo_tensor()
        x = self.mlp_adj(x)
        return x
        
    def reset_parameters(self):
        super().reset_parameters()
        self.mlp_adj.reset_parameters()
        if self.bn:
            self.bn.reset_parameters()

core/modules/node/em.py
thupchnsky/dp-gnn
from typing import Callable, Optional
import torch
from torch import Tensor
import torch.nn.functional as F
from core.models import MLP
from torch_geometric.data import Data
from core.modules.base import TrainableModule, Stage, Metrics
from torch_sparse import SparseTensor


class GraphNodeClassifier(TrainableModule):
    '''
        This is almost the same as MLPNodeClassifier, except that the input for the forward function
        is a bit different.
    '''
    def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,  
                 num_layers: int = 2, 
                 dropout: float = 0.0, 
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 ):

        super().__init__()

        self.model = MLP(
            output_dim=num_classes,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
            plain_last=True,
        )

    def forward(self, x: Tensor, adj_t: SparseTensor, noisy_x_emb: Tensor) -> Tensor:
        return self.model(x)

    def step(self, data: Data, stage: Stage) -> tuple[Optional[Tensor], Metrics]:
        mask = data[f'{stage}_mask']
        x, y, adj_t, noisy_x_emb = data.x, data.y[mask], data.adj_t, data.noisy_x_emb
        preds = F.log_softmax(self(x, adj_t, noisy_x_emb), dim=-1)[mask]
        acc = preds.argmax(dim=1).eq(y).float().mean() * 100
        metrics = {'acc': acc}

        loss = None
        if stage != 'test':
            loss = F.nll_loss(input=preds, target=y)
            metrics['loss'] = loss.detach()

        return loss, metrics

    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        h = self(data.x, data.adj_t, data.noisy_x_emb)
        return torch.softmax(h, dim=-1)

    def reset_parameters(self):
        return self.model.reset_parameters()
core/modules/node/graph_mlp.py
thupchnsky/dp-gnn
from typing import Callable, Optional
import torch
from torch import Tensor
import torch.nn.functional as F
from core.models import MLP
from torch_geometric.data import Data
from core.modules.base import TrainableModule, Stage, Metrics


class MLPNodeClassifier(TrainableModule):
    def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,  
                 num_layers: int = 2, 
                 dropout: float = 0.0, 
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 ):

        super().__init__()

        self.model = MLP(
            output_dim=num_classes,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
            plain_last=True,
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.model(x)

    def step(self, data: Data, stage: Stage) -> tuple[Optional[Tensor], Metrics]:
        mask = data[f'{stage}_mask']
        x, y = data.x[mask], data.y[mask]
        preds = F.log_softmax(self(x), dim=-1)
        acc = preds.argmax(dim=1).eq(y).float().mean() * 100
        metrics = {'acc': acc}

        loss = None
        if stage != 'test':
            loss = F.nll_loss(input=preds, target=y)
            metrics['loss'] = loss.detach()

        return loss, metrics

    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        h = self(data.x)
        return torch.softmax(h, dim=-1)

    def reset_parameters(self):
        return self.model.reset_parameters()


class MLPClassifier(TrainableModule):
    def __init__(self, *,
                 num_classes: int,
                 hidden_dim: int = 16,  
                 num_layers: int = 2, 
                 dropout: float = 0.0, 
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_, 
                 batch_norm: bool = False,
                 ):

        super().__init__()

        self.model = MLP(
            output_dim=num_classes,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
            plain_last=True,
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.model(x)

    def step(self, data: Data, stage: Stage) -> tuple[Optional[Tensor], Metrics]:
        x, y = data[0], data[1]
        preds = F.log_softmax(self(x), dim=-1)
        acc = preds.argmax(dim=1).eq(y).float().mean() * 100
        metrics = {'acc': acc}

        loss = None
        if stage != 'test':
            loss = F.nll_loss(input=preds, target=y)
            metrics['loss'] = loss.detach()

        return loss, metrics

    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        h = self(data[0])
        return torch.softmax(h, dim=-1)

    def reset_parameters(self):
        return self.model.reset_parameters()
core/modules/node/mlp.py
thupchnsky/dp-gnn
from typing import Callable
import torch
from torch import Tensor
import torch.nn.functional as F
from torch_sparse import SparseTensor
from core.models import MLP
from torch_geometric.data import Data
from core.models.sage import SAGE
from core.modules.base import Metrics, Stage, TrainableModule


class SAGENodeClassifier(TrainableModule):
    def __init__(self, *,
                 num_classes: int, 
                 hidden_dim: int = 16, 
                 base_layers: int = 0, 
                 mp_layers: int = 2, 
                 head_layers: int = 0, 
                 normalize: bool = False,
                 activation_fn: Callable[[Tensor], Tensor] = torch.relu_,
                 dropout: float = 0.0, 
                 batch_norm: bool = False,
                 ):

        assert mp_layers > 0, 'Must have at least one message passing layer'
        super().__init__()
        self.normalize = normalize

        self.base_mlp = MLP(
            hidden_dim=hidden_dim,
            output_dim=hidden_dim,
            num_layers=base_layers,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
            plain_last=False,
        )

        self.gnn = SAGE(
            output_dim=num_classes if head_layers == 0 else hidden_dim,
            hidden_dim=hidden_dim,
            num_layers=mp_layers,
            dropout=dropout,
            activation_fn=activation_fn,
            batch_norm=batch_norm,
            plain_last=head_layers == 0,
        )

        self.head_mlp = MLP(
            hidden_dim=hidden_dim,
            output_dim=num_classes,
            num_layers=head_layers,
            activation_fn=activation_fn,
            dropout=dropout,
            batch_norm=batch_norm,
            plain_last=True,
        )

    def forward(self, x: Tensor, adj_t: SparseTensor) -> Tensor:
        x = self.base_mlp(x)
        x = F.normalize(x, p=2, dim=-1) if self.normalize else x
        x = self.gnn(x, adj_t)
        x = self.head_mlp(x)
        return x

    def step(self, data: Data, stage: Stage) -> tuple[Tensor, Metrics]:
        mask = data[f'{stage}_mask']
        h = self(data.x, data.adj_t)[mask]
        preds = F.log_softmax(h, dim=-1)
        target = data.y[mask]
        acc = preds.argmax(dim=1).eq(target).float().mean() * 100
        metrics = {'acc': acc}

        loss = None
        if stage != 'test':
            loss = F.nll_loss(input=preds, target=target)
            metrics['loss'] = loss.detach()

        return loss, metrics

    @torch.no_grad()
    def predict(self, data: Data) -> Tensor:
        self.eval()
        h = self(data.x, data.adj_t)
        return torch.softmax(h, dim=-1)

    def reset_parameters(self):
        self.base_mlp.reset_parameters()
        self.gnn.reset_parameters()
        self.head_mlp.reset_parameters()

core/modules/node/sage.py
thupchnsky/dp-gnn
from core.privacy.algorithms.noisy_sgd import NoisySGD
from core.privacy.algorithms.graph import *

core/privacy/algorithms/__init__.py
thupchnsky/dp-gnn
from core.privacy.algorithms.graph.arr import AsymmetricRandResponse
from core.privacy.algorithms.graph.gnn_sgd import GNNBasedNoisySGD
from core.privacy.algorithms.graph.pma import PMA

core/privacy/algorithms/graph/__init__.py
thupchnsky/dp-gnn
import torch
import numpy as np
from torch_sparse import SparseTensor


class AsymmetricRandResponse:
    def __init__(self, eps: float):
        self.eps_link = eps * 0.9
        self.eps_density = eps * 0.1

    def __call__(self, adj_t: SparseTensor, chunk_size: int=1000) -> SparseTensor:
        chunks = self.split(adj_t, chunk_size=chunk_size)
        pert_chunks = []

        for chunk in chunks:    
            chunk_pert = self.perturb(chunk)
            pert_chunks.append(chunk_pert)

        perturbed_adj_t = self.merge(pert_chunks, chunk_size=chunk_size)
        return perturbed_adj_t
    
    def split(self, adj_t: SparseTensor, chunk_size: int) -> list[SparseTensor]:
        chunks = []
        for i in range(0, adj_t.size(0), chunk_size):
            if (i + chunk_size) <= adj_t.size(0):
                chunks.append(adj_t[i:i+chunk_size])
            else:
                chunks.append(adj_t[i:])
        return chunks
    
    def perturb(self, adj_t: SparseTensor) -> SparseTensor:
        n = adj_t.size(1)
        sensitivity = 1 / (n*n)
        p = 1 / (1 + np.exp(-self.eps_link))
        d = np.random.laplace(loc=adj_t.density(), scale=sensitivity/self.eps_density)
        q = d / (2*p*d - p - d + 1)
        q = min(1, q)
        pr_1to1 = p * q
        pr_0to1 = (1 - p) * q
        mask = adj_t.to_dense(dtype=bool)
        out = mask * pr_1to1 + (~mask) * pr_0to1
        torch.bernoulli(out, out=out)
        out = SparseTensor.from_dense(out, has_value=False)
        return out
    
    def merge(self, chunks: list[SparseTensor], chunk_size: int) -> SparseTensor:
        n = (len(chunks) - 1) * chunk_size + chunks[-1].size(0)
        m = chunks[0].size(1)
        row = torch.cat([chunk.coo()[0] + i * chunk_size for i, chunk in enumerate(chunks)])
        col = torch.cat([chunk.coo()[1] for chunk in chunks])
        out = SparseTensor(row=row, col=col, sparse_sizes=(n, m))#.coalesce()
        return out

core/privacy/algorithms/graph/arr.py
thupchnsky/dp-gnn
from typing import TypeVar
import numpy as np
from scipy.stats import hypergeom
from opacus.grad_sample import GradSampleModule
from opacus.optimizers import DPOptimizer
from opacus.privacy_engine import forbid_accumulation_hook
from torch.nn import Module
from torch.optim import Optimizer
from torch.utils.data import DataLoader
from core.privacy.mechanisms.commons import InfMechanism, ZeroMechanism
from core.privacy.mechanisms.noisy import NoisyMechanism


T = TypeVar('T', bound=Module)


class GNNBasedNoisySGD(NoisyMechanism):
    def __init__(self, noise_scale: float, dataset_size: int, batch_size: int, 
                 epochs: int, max_grad_norm: float, max_degree: int):
        super().__init__(noise_scale)
        self.name = 'NoisySGD'
        self.params = {
            'noise_scale': noise_scale, 
            'dataset_size': dataset_size, 
            'batch_size': batch_size, 
            'epochs': epochs,
            'max_grad_norm': max_grad_norm,
            'max_degree': max_degree,
        }

        if epochs == 0:
            mech = ZeroMechanism()
            self.params['noise_scale'] = 0.0
            self.set_all_representation(mech)
        elif noise_scale == 0.0:
            mech = InfMechanism()
            self.set_all_representation(mech)
        else:
            N = dataset_size
            K = max_degree
            m = batch_size
            C = max_grad_norm
            T = epochs * dataset_size // batch_size
            DeltaK = 2 * (K + 1) * C
            
            def RDP(alpha):
                sigma = noise_scale * DeltaK
                expected_fn = lambda rho: np.exp(alpha * (alpha-1) * 2 * rho**2 * C**2 / sigma**2)
                expected_rho = hypergeom(N, K+1, m).expect(expected_fn)
                gamma = np.log(expected_rho) / (alpha - 1)
                return gamma * T

            self.propagate_updates(RDP, type_of_update='RDP')

    def __call__(self, module: T, optimizer: Optimizer, dataloader: DataLoader) -> tuple[T, Optimizer, DataLoader]:
        module = self.prepare_module(module)
        dataloader = self.prepare_dataloader(dataloader)
        optimizer = self.prepare_optimizer(optimizer)
        return module, optimizer, dataloader

    def prepare_module(self, module: T) -> T:
        if self.params['noise_scale'] > 0.0 and self.params['epochs'] > 0:
            if hasattr(module, 'autograd_grad_sample_hooks'):
                for hook in module.autograd_grad_sample_hooks:
                    hook.remove()
                del module.autograd_grad_sample_hooks
            GradSampleModule(module).register_backward_hook(forbid_accumulation_hook)
        return module

    def prepare_dataloader(self, dataloader: DataLoader) -> DataLoader:
        # since we don't need poisson sampling, we can use the same dataloader
        return dataloader

    def prepare_optimizer(self, optimizer: Optimizer) -> DPOptimizer:
        noise_scale = self.params['noise_scale']
        epochs = self.params['epochs']
        K = self.params['max_degree']
        C = self.params['max_grad_norm']

        if noise_scale > 0.0 and epochs > 0:
            optimizer = DPOptimizer(
                optimizer=optimizer,
                noise_multiplier=noise_scale * 2 * (K + 1),
                max_grad_norm=C,
                expected_batch_size=self.params['batch_size'],
            )
        return optimizer

core/privacy/algorithms/graph/gnn_sgd.py
thupchnsky/dp-gnn
import torch
from core.privacy.mechanisms.commons import GaussianMechanism, InfMechanism, ZeroMechanism
from core.privacy.mechanisms.composed import ComposedGaussianMechanism
from core.privacy.mechanisms.noisy import NoisyMechanism


class PMA(NoisyMechanism):
    def __init__(self, noise_scale: float, hops: int):
        super().__init__(noise_scale)
        self.name = 'PMA'
        self.params = {'noise_scale': noise_scale, 'hops': hops}
        self.gm = GaussianMechanism(noise_scale=noise_scale) 

        if hops == 0:
            mech = ZeroMechanism()
            self.params['noise_scale'] = 0.0
        elif noise_scale == 0.0:
            mech = InfMechanism()
        else:
            mech = ComposedGaussianMechanism(
                noise_scale=noise_scale, 
                mechanism_list=[self.gm], 
                coeff_list=[hops]
            )

        self.set_all_representation(mech)

    def __call__(self, x: torch.Tensor, sensitivity: float) -> torch.Tensor:
        return self.gm(x, sensitivity=sensitivity)

core/privacy/algorithms/graph/pma.py
thupchnsky/dp-gnn
from typing import TypeVar
from opacus.privacy_engine import forbid_accumulation_hook
from opacus.grad_sample import GradSampleModule
from opacus.optimizers import DPOptimizer
from autodp.transformer_zoo import Composition, AmplificationBySampling
from torch.nn import Module
from torch.optim import Optimizer
from core.data.loader import NodeDataLoader
from core.privacy.mechanisms.commons import GaussianMechanism, InfMechanism, ZeroMechanism
from core.privacy.mechanisms.noisy import NoisyMechanism

T = TypeVar('T', bound=Module)

class NoisySGD(NoisyMechanism):
    def __init__(self, noise_scale: float,
                 dataset_size: int,
                 batch_size: int,
                 epochs: int,
                 max_grad_norm: float,
                 group_size: int=1,
                 replacement: bool=False):
        super().__init__(noise_scale)
        self.name = 'NoisySGD'
        self.params = {
            'noise_scale': noise_scale, 
            'dataset_size': dataset_size, 
            'batch_size': batch_size, 
            'epochs': epochs,
            'max_grad_norm': max_grad_norm,
            'group_size': group_size,
            'replacement': replacement
        }
        
        if epochs == 0:
            mech = ZeroMechanism()
            self.params['noise_scale'] = 0.0
        elif noise_scale == 0.0:
            mech = InfMechanism()
        else:
            if self.params['replacement']:
                subsample = AmplificationBySampling(PoissonSampling=False)
            else:
                subsample = AmplificationBySampling(PoissonSampling=True)
            compose = Composition()
            gm = GaussianMechanism(noise_scale=noise_scale,replacement=self.params['replacement'])
            if self.params['replacement']:
                gm.neighboring = 'replace_one'
                gm.replace_one = True
            subsampled_gm = subsample(gm, prob=batch_size/dataset_size, improved_bound_flag=True)
            mech = compose([subsampled_gm],[epochs * dataset_size // batch_size])
        
        
        self.set_all_representation(mech)
        
        
    def __call__(self, module: T, optimizer: Optimizer, dataloader: NodeDataLoader) -> tuple[T, Optimizer, NodeDataLoader]:
        module = self.prepare_module(module)
        dataloader = self.prepare_dataloader(dataloader)
        optimizer = self.prepare_optimizer(optimizer)
        return module, optimizer, dataloader

    def prepare_module(self, module: T) -> T:
        if self.params['noise_scale'] > 0.0 and self.params['epochs'] > 0:
            if hasattr(module, 'autograd_grad_sample_hooks'):
                for hook in module.autograd_grad_sample_hooks:
                    hook.remove()
                del module.autograd_grad_sample_hooks
            GradSampleModule(module).register_backward_hook(forbid_accumulation_hook)
        return module

    def prepare_dataloader(self, dataloader: NodeDataLoader) -> NodeDataLoader:
        if self.params['noise_scale'] > 0.0 and self.params['epochs'] > 0:
            if self.replacement:
                dataloader.wor_sampling = True
            else:
                dataloader.poisson_sampling = True
        return dataloader

    def prepare_optimizer(self, optimizer: Optimizer, noise_factor: float = 1.0) -> DPOptimizer:
        if self.params['noise_scale'] > 0.0 and self.params['epochs'] > 0:
            optimizer = DPOptimizer(
                optimizer=optimizer,
                noise_multiplier=self.params['noise_scale']*noise_factor,    # noise_multiplier is the same as noise_scale in Opacus
                max_grad_norm=self.params['max_grad_norm'],
                expected_batch_size=self.params['batch_size'],
            )
        return optimizer

core/privacy/algorithms/noisy_sgd.py
thupchnsky/dp-gnn
from core.privacy.mechanisms.commons import ZeroMechanism, InfMechanism, GaussianMechanism, LaplaceMechanism
from core.privacy.mechanisms.composed import ComposedNoisyMechanism, ComposedGaussianMechanism
from core.privacy.mechanisms.noisy import NoisyMechanism

core/privacy/mechanisms/__init__.py
thupchnsky/dp-gnn
import torch
from torch import Tensor
from autodp.mechanism_zoo import Mechanism, ExactGaussianMechanism, LaplaceMechanism as AutoDPLaplaceMechanism
from core.privacy.mechanisms.noisy import NoisyMechanism


class ZeroMechanism(Mechanism):
    def __init__(self):
        super().__init__()
        self.name = 'ZeroMechanism'
        self.params = {}
        self.propagate_updates(func=lambda _: 0, type_of_update='RDP')
        

class InfMechanism(Mechanism):
    def __init__(self):
        super().__init__()
        self.name = 'InfMechanism'
        self.params = {}


class GaussianMechanism(NoisyMechanism):
    def __init__(self, noise_scale: float,
                 group_size: int=1,
                 replacement: bool=False):
        # "noise_scale" is the std of the noise divide by the L2 sensitivity
        super().__init__(noise_scale=noise_scale)
        gm = ExactGaussianMechanism(sigma=noise_scale)
        gm.group_size = group_size
        if replacement:
            gm.neighboring = 'replace_one'
            gm.replace_one = True

        self.name = 'GaussianMechanism'
        self.set_all_representation(gm)

    def perturb(self, data: Tensor, sensitivity: float) -> Tensor:
        std = self.params['noise_scale'] * sensitivity
        return torch.normal(mean=data, std=std) if std else data

    def __call__(self, data: Tensor, sensitivity: float) -> Tensor:
        return self.perturb(data, sensitivity)


class LaplaceMechanism(NoisyMechanism):
    def __init__(self, noise_scale: float):
        # "noise_scale" is the Laplace scale parameter divided by the L1 sensitivity
        super().__init__(noise_scale=noise_scale)
        lm = AutoDPLaplaceMechanism(b=noise_scale)
        self.name = 'LaplaceMechanism'
        self.set_all_representation(lm)

    def perturb(self, data: Tensor, sensitivity: float) -> Tensor:
        scale = self.params['noise_scale'] * sensitivity
        return torch.distributions.Laplace(loc=data, scale=scale).sample() if scale else data

    def __call__(self, data: Tensor, sensitivity: float) -> Tensor:
        return self.perturb(data, sensitivity)

core/privacy/mechanisms/commons.py
thupchnsky/dp-gnn
from core.privacy.mechanisms.commons import GaussianMechanism
from core.privacy.mechanisms.noisy import NoisyMechanism
from autodp.transformer_zoo import ComposeGaussian, Composition
from autodp.mechanism_zoo import ExactGaussianMechanism


class ComposedNoisyMechanism(NoisyMechanism):
    def __init__(self, 
                 noise_scale: float, 
                 mechanism_list: list[NoisyMechanism], 
                 coeff_list: list[float]=None,
                 weight_list: list[float]=None,
                 ):
        super().__init__(noise_scale)
        if coeff_list is None:
            coeff_list = [1] * len(mechanism_list)
        if weight_list is None:
            weight_list = [1] * len(mechanism_list)
        self.params = {
            'noise_scale': noise_scale, 
            'mechanism_list': mechanism_list, 
            'coeff_list': coeff_list,
            'weight_list': weight_list
        }
        mechanism_list = [mech.update(weight * noise_scale) for mech, weight in zip(mechanism_list, weight_list)]
        mech = Composition()(mechanism_list, coeff_list)
        self.set_all_representation(mech)


class ComposedGaussianMechanism(NoisyMechanism):
    def __init__(self, 
                 noise_scale: float, 
                 mechanism_list: list[GaussianMechanism], 
                 coeff_list: list[float]
                 ):
        super().__init__(noise_scale)
        self.params = {'noise_scale': noise_scale, 'mechanism_list': mechanism_list, 'coeff_list': coeff_list}
        mechanism_list = [mech.update(noise_scale) for mech in mechanism_list]
        gm_list = [ExactGaussianMechanism(sigma=mech.params['noise_scale']) for mech in mechanism_list]
        mech = ComposeGaussian()(gm_list, coeff_list)
        self.set_all_representation(mech)
core/privacy/mechanisms/composed.py
thupchnsky/dp-gnn
from typing_extensions import Self
from scipy.optimize import minimize_scalar, OptimizeResult
from autodp.mechanism_zoo import Mechanism
import numpy as np
from core import console


class NoisyMechanism(Mechanism):
    def __init__(self, noise_scale: float):
        # "noise_scale" is the std of the noise divide by the sensitivity
        super().__init__()
        self.name = 'NoisyMechanism'
        self.params = {'noise_scale': noise_scale}

    def update(self, noise_scale: float) -> Self:
        self.params.pop('noise_scale')
        self.__init__(noise_scale, **self.params)
        return self

    def is_zero(self) -> bool:
        for alpha in range(2,100):
            if self.RenyiDP(alpha) > 0:
                return False
        return True

    def is_inf(self) -> bool:
        for alpha in range(2,100):
            if not np.isinf(self.RenyiDP(alpha)):
                return False
        return True

    def calibrate(self, eps: float, delta: float) -> float:
        if self.params['noise_scale'] == 0:
            self.update(noise_scale=1)  # to avoid is_inf being true
        
        console.debug('checking if the mechanism is inf or zero...')
        if np.isinf(eps) or self.is_inf() or self.is_zero():
            self.update(noise_scale=0)
            return 0.0
        else:
            console.debug('calibration begins...')
            fn_err = lambda x: abs(eps - self.update(np.exp(x)).get_approxDP(delta))

            # check if the mechanism is already calibrated
            if fn_err(self.params['noise_scale']) < 1e-3:
                return self.params['noise_scale']
            
            result: OptimizeResult = minimize_scalar(fn_err, 
                method='brent', 
                options={'xtol': 1e-5, 'maxiter': 1000000, 'disp': 0}
            )

            if result.success and result.fun < 1e-3:
                noise_scale = np.exp(result.x)
                self.update(noise_scale)
                return noise_scale
            else:
                raise RuntimeError(f"calibrator failed to find noise scale\n{result}")
core/privacy/mechanisms/noisy.py
thupchnsky/dp-gnn
from core.trainer.trainer import Trainer

core/trainer/__init__.py
thupchnsky/dp-gnn
from typing import Iterable
from rich.console import Group
from rich.padding import Padding
from rich.text import Text
from rich.table import Column, Table
from core.modules.base import Metrics
from core import console
from rich.progress import Progress, SpinnerColumn, BarColumn, TimeElapsedColumn, Task
from rich.highlighter import ReprHighlighter


class TrainerProgress(Progress):
    def __init__(self, 
                 num_epochs: int, 
                 num_train_steps: int, 
                 num_val_steps: int, 
                 num_test_steps: int,
                 ):

        progress_bar = [
            SpinnerColumn(),
            "{task.description}",
            "[cyan]{task.completed:>3}[/cyan]/[cyan]{task.total}[/cyan]",
            "{task.fields[unit]}",
            BarColumn(),
            "[cyan]{task.percentage:>3.0f}[/cyan]%",
            TimeElapsedColumn(),
            # "{task.fields[metrics]}"
        ]

        super().__init__(*progress_bar, console=console)

        self.trainer_tasks = {
            'epoch': self.add_task(total=num_epochs, metrics='', unit='epochs', description='overal progress'),
            'train': self.add_task(total=num_train_steps, metrics='', unit='steps', description='training', visible=False),
            'val':   self.add_task(total=num_val_steps, metrics='', unit='steps', description='validation', visible=False),
            'test':  self.add_task(total=num_test_steps, metrics='', unit='steps', description='testing', visible=False),
        }

        self.max_rows = 0

    def update(self, task: Task, **kwargs):
        if 'metrics' in kwargs:
            kwargs['metrics'] = self.render_metrics(kwargs['metrics'])

        super().update(self.trainer_tasks[task], **kwargs)

    def reset(self, task: Task, **kwargs):
        super().reset(self.trainer_tasks[task], **kwargs)

    def render_metrics(self, metrics: Metrics) -> str:
        out = []
        for split in ['train', 'val', 'test']:
            metric_str = ' '.join(f'{k}: {v:.3f}' for k, v in metrics.items() if f'{split}/' in k)
            out.append(metric_str)
        
        return '  '.join(out)

    def make_tasks_table(self, tasks: Iterable[Task]) -> Table:
        """Get a table to render the Progress display.

        Args:
            tasks (Iterable[Task]): An iterable of Task instances, one per row of the table.

        Returns:
            Table: A table instance.
        """
        table_columns = (
            (
                Column(no_wrap=True)
                if isinstance(_column, str)
                else _column.get_table_column().copy()
            )
            for _column in self.columns
        )

        highlighter = ReprHighlighter()
        table = Table.grid(*table_columns, padding=(0, 1), expand=self.expand)

        if tasks:
            epoch_task = tasks[0]
            metrics = epoch_task.fields['metrics']

            for task in tasks:
                if task.visible:
                    table.add_row(
                        *(
                            (
                                column.format(task=task)
                                if isinstance(column, str)
                                else column(task)
                            )
                            for column in self.columns
                        )
                    )

            self.max_rows = max(self.max_rows, table.row_count)
            pad_top = 0 if epoch_task.finished else self.max_rows - table.row_count
            group = Group(table, Padding(highlighter(metrics), pad=(pad_top,0,0,2)))
            return Padding(group, pad=(0,0,1,21))

        else:
            return table
core/trainer/progress.py
thupchnsky/dp-gnn
import os
import uuid
import torch
from torch.types import Number
from torch.optim import Optimizer
from typing import Annotated, Iterable, Literal, Optional
from core.args.utils import ArgInfo
from core.loggers import Logger
from torchmetrics import MeanMetric
from core.trainer.progress import TrainerProgress
from core.modules.base import Metrics, Stage, TrainableModule


class Trainer:
    def __init__(self,
                 patience:      int = 0,
                 monitor:       str = 'val/acc',
                 monitor_mode:  Literal['min', 'max'] = 'max',
                 val_interval:  Annotated[int, ArgInfo(help='interval of validation')] = 1,
                 ):

        assert monitor_mode in ['min', 'max']

        self.patience = patience
        self.val_interval = val_interval
        self.monitor = monitor
        self.monitor_mode = monitor_mode
        
        # trainer internal state
        self.model: TrainableModule = None
        self.metrics: dict[str, MeanMetric] = {}

    def reset(self):
        self.model = None
        self.metrics = {}

    def update_metrics(self, metric_name: str, metric_value: object, weight: int = 1) -> None:
        # if this is a new metric, add it to self.metrics
        device = metric_value.device if torch.is_tensor(metric_value) else 'cpu'
        if metric_name not in self.metrics:
            self.metrics[metric_name] = MeanMetric(compute_on_step=False).to(device)

        # update the metric
        self.metrics[metric_name].update(metric_value, weight=weight)

    def aggregate_metrics(self, stage: Stage='train') -> Metrics:
        metrics = {}

        for metric_name, metric_value in self.metrics.items():
            if stage in metric_name.split('/'):
                value = metric_value.compute()
                metric_value.reset()
                if torch.is_tensor(value):
                    value = value.item()
                metrics[metric_name] = value

        return metrics

    def is_better(self, current_metric: Number, previous_metric: Number) -> bool:
        if self.monitor_mode == 'max':
            return current_metric > previous_metric
        elif self.monitor_mode == 'min':
            return current_metric < previous_metric
        else:
            raise ValueError(f'Unknown metric mode: {self.monitor_mode}')

    def fit(self, 
            model: TrainableModule, 
            epochs: int,
            optimizer: Optimizer, 
            train_dataloader: Iterable, 
            val_dataloader: Optional[Iterable]=None, 
            test_dataloader: Optional[Iterable]=None, 
            checkpoint: bool=False,
            prefix: str = ''
            ) -> Metrics:

        self.model = model
        self.optimizer = optimizer
        monitor_key = f'{prefix}{self.monitor}'

        if checkpoint:
            os.makedirs('checkpoints', exist_ok=True)
            checkpoint_path = os.path.join('checkpoints', f'{uuid.uuid1()}.pt')
            torch.save(self.model.state_dict(), checkpoint_path)

        if val_dataloader is None:
            val_dataloader = []

        if test_dataloader is None:
            test_dataloader = []

        self.progress = TrainerProgress(
            num_epochs=epochs, 
            num_train_steps=len(train_dataloader), 
            num_val_steps=len(val_dataloader), 
            num_test_steps=len(test_dataloader),
        )
        
        with self.progress:
            best_metrics = None
            num_epochs_without_improvement = 0
            
            for epoch in range(1, epochs + 1):
                metrics = {f'{prefix}epoch': epoch}

                # train loop
                train_metrics = self.loop(train_dataloader, stage='train', prefix=prefix)
                metrics.update(train_metrics)
                    
                # validation loop
                if val_dataloader and self.val_interval and epoch % self.val_interval == 0:

                    val_metrics = self.loop(val_dataloader, stage='val', prefix=prefix)
                    metrics.update(val_metrics)

                    if best_metrics is None or self.is_better(metrics[monitor_key], best_metrics[monitor_key]):
                        best_metrics = metrics
                        num_epochs_without_improvement = 0

                        if checkpoint:
                            torch.save(self.model.state_dict(), checkpoint_path)
                    else:
                        num_epochs_without_improvement += 1
                        if num_epochs_without_improvement >= self.patience > 0:
                            break

                # test loop
                if test_dataloader:
                    test_metrics = self.loop(test_dataloader, stage='test', prefix=prefix)
                    metrics.update(test_metrics)

                # log and update progress
                Logger.get_instance().log(metrics)
                self.progress.update(task='epoch', metrics=metrics, advance=1)

        if best_metrics is None:
            best_metrics = metrics
        else:
            # load best model if checkpointing is enabled
            if checkpoint:
                self.model.load_state_dict(torch.load(checkpoint_path))

        # log and return best metrics
        Logger.get_instance().log_summary(best_metrics)
        return best_metrics

    def test(self, dataloader: Iterable, prefix: str = '') -> Metrics:
        metrics = self.loop(dataloader, stage='test', prefix=prefix)
        return metrics

    def loop(self, dataloader: Iterable, stage: Stage, prefix: str) -> Metrics:
        self.model.train(stage == 'train')
        self.progress.update(stage, visible=len(dataloader) > 1)

        for batch in dataloader:
            metrics = self.step(batch, stage, prefix)
            for item in metrics:
                self.update_metrics(item, metrics[item], weight=len(batch))
            self.progress.update(stage, advance=1)

        self.progress.reset(stage, visible=False)
        return self.aggregate_metrics(stage)

    def step(self, batch, stage: Stage, prefix: str) -> Metrics:
        if stage == 'train':
            self.optimizer.zero_grad(set_to_none=True)

        grad_state = torch.is_grad_enabled()
        torch.set_grad_enabled(stage == 'train')
        loss, metrics = self.model.step(batch, stage=stage)
        torch.set_grad_enabled(grad_state)
        
        if stage == 'train' and loss is not None:
            loss.backward()
            self.optimizer.step()

        return {f'{prefix}{stage}/{key}': value for key, value in metrics.items()}

core/trainer/trainer.py
thupchnsky/dp-gnn
import math
import random
from typing import Callable, Iterable, TypeVar, Optional
import numpy as np
import torch
from itertools import tee
from numpy.typing import ArrayLike, NDArray
from rich.table import Table
from rich.highlighter import ReprHighlighter
from rich import box
from tabulate import tabulate


RT = TypeVar('RT')


def seed_everything(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def pairwise(iterable: Iterable):
    # pairwise('ABCDEFG') --> AB BC CD DE EF FG
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


def bootstrap(data: ArrayLike, 
              func: Callable[[ArrayLike], NDArray]=np.mean, 
              n_boot: int=10000, 
              seed: Optional[int]=None
              ) -> NDArray:
    n = len(data)
    data = np.asarray(data)
    rng = np.random.default_rng(seed)
    integers = rng.integers
    
    boot_dist = []
    for _ in range(int(n_boot)):
        resampler = integers(0, n, n, dtype=np.intp)  # intp is indexing dtype
        sample = [data.take(resampler, axis=0)]
        boot_dist.append(func(*sample))
        
    return np.array(boot_dist)


def confidence_interval(data: ArrayLike, 
                        func: Callable[[ArrayLike], NDArray]=np.mean, 
                        size: int=1000, 
                        ci: int=95, 
                        seed: Optional[int]=None
                        ) -> float:
    bs_replicates = bootstrap(data, func=func, n_boot=size, seed=seed)
    p = 50 - ci / 2, 50 + ci / 2
    bounds = np.nanpercentile(bs_replicates, p)
    return (bounds[1] - bounds[0]) / 2


def dict2table(input_dict: dict, num_cols: int = 4, title: Optional[str] = None) -> Table:
    num_items = len(input_dict)
    num_rows = math.ceil(num_items / num_cols)
    col = 0
    data = {}
    keys = []
    vals = []

    for i, (key, val) in enumerate(input_dict.items()):
        keys.append(f'{key}:')
        
        vals.append(val)
        if (i + 1) % num_rows == 0:
            data[col] = keys
            data[col+1] = vals
            keys = []
            vals = []
            col += 2

    data[col] = keys
    data[col+1] = vals

    highlighter = ReprHighlighter()
    message = tabulate(data, tablefmt='plain')
    table = Table(title=title, show_header=False, box=box.HORIZONTALS)
    table.add_row(highlighter(message))
    return table

core/utils.py
thupchnsky/dp-gnn
from core import console
with console.status('importing modules'):
    import torch
    import numpy as np
    from rich import box
    from rich.table import Table
    from time import time
    from typing import Annotated
    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
    from core import globals
    from core.datasets import DatasetLoader
    from core.args.utils import print_args, create_arguments, strip_kwargs, ArgInfo
    from core.loggers import Logger
    from core.methods.node import supported_methods, NodeClassification
    from core.utils import seed_everything, confidence_interval
    from torch_geometric.data import Data


def run(seed:     Annotated[int,   ArgInfo(help='initial random seed')] = 12345,
        repeats:  Annotated[int,   ArgInfo(help='number of times the experiment is repeated')] = 1,
        csbm_phi: Annotated[float,   ArgInfo(help='parameter phi for cSBM dataset generation')] = 0.0,
        debug:    Annotated[bool,  ArgInfo(help='enable global debug mode')] = False,
        **kwargs
    ):

    seed_everything(seed)

    if debug:
        console.info('debug mode enabled')
        globals['debug'] = True
        console.log_level = console.DEBUG

    with console.status('loading dataset'):
        loader_args = strip_kwargs(DatasetLoader, kwargs)
        data_initial = DatasetLoader(**loader_args).load(verbose=True, csbm_phi=csbm_phi)

    num_classes = data_initial.y.max().item() + 1
    num_features = data_initial.num_features
    config = dict(**kwargs, seed=seed, repeats=repeats)
    logger_args = strip_kwargs(Logger.setup, kwargs)
    logger = Logger.setup(enabled=False, config=config, **logger_args)

    ### initiallize method ###
    Method = supported_methods[kwargs['method']]
    method_args = strip_kwargs(Method, kwargs)
    if kwargs['method'].startswith('dpdgc'):
        method = Method(num_features=num_features, num_classes=num_classes, **method_args)
    else:
        method = Method(num_classes=num_classes, **method_args)

    run_metrics = {}

    ### run experiment ###
    for iteration in range(repeats):
        start_time = time()
        data = Data(**data_initial.to_dict())
        metrics = method.fit(data)
        end_time = time()
        metrics['duration'] = end_time - start_time

        ### process results ###
        for metric, value in metrics.items():
            run_metrics[metric] = run_metrics.get(metric, []) + [value]

         ### print results ###
        table = Table(title=f'run {iteration + 1}', box=box.HORIZONTALS)
        table.add_column('metric')
        table.add_column('last', style="cyan")
        table.add_column('mean', style="cyan")
        table.add_row('test/acc', f'{run_metrics["test/acc"][-1]:.2f}', f'{np.mean(run_metrics["test/acc"]):.2f}')
        console.info(table)
        console.print()

        ### reset method's parameters for the next run ###
        method.reset_parameters()

    logger.enable()
    summary = {}
    
    for metric, values in run_metrics.items():
        summary[metric + '_mean'] = np.mean(values)
        summary[metric + '_std'] = np.std(values)
        summary[metric + '_ci'] = confidence_interval(values, size=1000, ci=95, seed=seed)
        logger.log_summary(summary)

    logger.finish()
    print()


def main():
    init_parser = ArgumentParser(add_help=False, conflict_handler='resolve')
    method_subparser = init_parser.add_subparsers(dest='method', required=True, title='algorithm to use')

    for method_name, method_class in supported_methods.items():
        method_parser = method_subparser.add_parser(
            name=method_name, 
            help=method_class.__doc__, 
            formatter_class=ArgumentDefaultsHelpFormatter
        )

        # dataset args
        group_dataset = method_parser.add_argument_group('dataset arguments')
        create_arguments(DatasetLoader, group_dataset)

        # method args
        group_method = method_parser.add_argument_group('method arguments')
        create_arguments(method_class, group_method)
        
        # experiment args
        group_expr = method_parser.add_argument_group('experiment arguments')
        create_arguments(run, group_expr)
        create_arguments(Logger.setup, group_expr)

    parser = ArgumentParser(parents=[init_parser], formatter_class=ArgumentDefaultsHelpFormatter)
    kwargs = vars(parser.parse_args())
    print_args(kwargs, num_cols=2)

    try:
        start = time()
        run(**kwargs)
        end = time()
        console.info(f'Total running time: {(end - start):.2f} seconds.')
    except KeyboardInterrupt:
        print('\n')
        console.warning('Graceful Shutdown')
    except RuntimeError:
        raise
    finally:
        if torch.cuda.is_available():
            gpu_mem = torch.cuda.max_memory_allocated() / 1024 ** 3
            console.info(f'Max GPU memory used = {gpu_mem:.2f} GB\n')


if __name__ == '__main__':
    main()

train.py
thupchnsky/dp-gnn