Official repository for "Activation Compression of Graph Neural Networks using Block-wise Quantization and Improved Variance Minimization". Eliassen and Selvan. 2023
from .utils import get_memory_usage, compute_tensor_bytes, exp_recorder, cast_adj, cast_low_bit_int
from .layers import QGCNConv, QLinear, QReLU, QBatchNorm1d, QDropout, QSAGEConv, QELU, QGCN2Conv, QGENConv
from .conf import config
from .module import QModule
from .gatconv import CustomGATConv, QCustomGATConv
exact/build/lib.linux-x86_64-cpython-37/exact/__init__.py
saintslab/i-Exact
class QuantizationConfig:
    def __init__(self):
        self.compress_activation = True
        self.activation_compression_bits = [2]
        self.initial_bits = 8
        self.stochastic = True
        self.use_gradient = False
        self.simulate = False
        self.enable_quantized_bn = True
        self.kept_frac = 1.0
        self.single_precision = True
        self.dropout2 = False
        self.col_size = 0
        self.last_epoch = False
        self.lo = 1
        self.use_optimal_lo = False
        # Memory management flag
        self.empty_cache_threshold = None
        self.swap = False
        self.amp = False
        # Debug related flag
        self.debug_memory_op_forward = False
        self.debug_memory_op_backward = False
        self.debug_remove_bn = False
        self.debug_remove_relu = False
        self.debug_remove_dropout = False

config = QuantizationConfig()
exact/build/lib.linux-x86_64-cpython-37/exact/conf.py
saintslab/i-Exact
from exact.layers import QDropout, QLinear
from typing import Union, Tuple, Optional
from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
                                    OptTensor)

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import Parameter, Linear
from torch_sparse import SparseTensor, set_diag
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops, softmax

from torch_geometric.nn.inits import zeros
import exact.cpp_extension.quantization as ext_quantization


class CustomGATConv(MessagePassing):
    def __init__(self, in_channels: Union[int, Tuple[int, int]],
                 out_channels: int, heads: int = 1, concat: bool = True,
                 negative_slope: float = 0.2, dropout: float = 0.0,
                 add_self_loops: bool = True, bias: bool = True, 
                 residual=True, use_attn_dst=True, **kwargs):
        kwargs.setdefault('aggr', 'add')
        super(CustomGATConv, self).__init__(node_dim=0, **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.negative_slope = negative_slope
        self.dropout = dropout
        self.add_self_loops = add_self_loops
        self.residual = residual
        self.use_attn_dst = use_attn_dst

        if isinstance(in_channels, int):
            self.lin_l = Linear(in_channels, heads * out_channels, bias=False)
            self.lin_r = self.lin_l
        else:
            self.lin_l = Linear(in_channels[0], heads * out_channels, False)
            self.lin_r = Linear(in_channels[1], heads * out_channels, False)

        if residual:
            self.res_fc = Linear(in_channels, heads * out_channels, bias=False)
        else:
            self.register_buffer("res_fc", None)

        self.att_l = Parameter(torch.Tensor(1, heads, out_channels))
        if self.use_attn_dst:
            self.att_r = Parameter(torch.Tensor(1, heads, out_channels))
        else:
            self.att_r = self.att_l
        if bias and concat:
            self.bias = Parameter(torch.Tensor(heads * out_channels))
        elif bias and not concat:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()


    def reset_parameters(self):
        gain = torch.nn.init.calculate_gain("relu")
        torch.nn.init.xavier_normal_(self.lin_l.weight, gain=gain)
        torch.nn.init.xavier_normal_(self.lin_r.weight, gain=gain)
        if isinstance(self.res_fc, Linear):
            torch.nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
        torch.nn.init.xavier_normal_(self.att_l, gain=gain)
        if self.use_attn_dst:
            torch.nn.init.xavier_normal_(self.att_r, gain=gain)
        zeros(self.bias)


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None):
        H, C = self.heads, self.out_channels

        x_l: OptTensor = None
        x_r: OptTensor = None
        alpha_l: OptTensor = None
        alpha_r: OptTensor = None
        if isinstance(x, Tensor):
            assert x.dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = x_r = self.lin_l(x).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if self.use_attn_dst:
                alpha_r = (x_r * self.att_r).sum(dim=-1)
            else:
                alpha_r = alpha_l
        else:
            x_l, x_r = x[0], x[1]
            assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = self.lin_l(x_l).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if x_r is not None:
                x_r = self.lin_r(x_r).view(-1, H, C)
                alpha_r = (x_r * self.att_r).sum(dim=-1)

        assert x_l is not None
        assert alpha_l is not None

        if self.add_self_loops:
            if isinstance(edge_index, Tensor):
                num_nodes = x_l.size(0)
                if x_r is not None:
                    num_nodes = min(num_nodes, x_r.size(0))
                if size is not None:
                    num_nodes = min(size[0], size[1])
                edge_index, _ = remove_self_loops(edge_index)
                edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
            elif isinstance(edge_index, SparseTensor):
                edge_index = set_diag(edge_index)

        # propagate_type: (x: OptPairTensor, alpha: OptPairTensor)
        out = self.propagate(edge_index, x=(x_l, x_r),
                             alpha=(alpha_l, alpha_r), size=size)


        if self.concat:
            out = out.view(-1, self.heads * self.out_channels)
        else:
            out = out.mean(dim=1)
        if self.residual:
            out += self.res_fc(x)
        if self.bias is not None:
            out += self.bias
        return out

    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
                index: Tensor, ptr: OptTensor,
                size_i: Optional[int]) -> Tensor:
        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
        alpha = F.leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
        return x_j * alpha.unsqueeze(-1)

    def __repr__(self):
        return '{}({}, {}, heads={}, residual={}, use_attn_dst={})'.format(self.__class__.__name__,
                                                                           self.in_channels, self.out_channels, 
                                                                           self.heads, self.residual, 
                                                                           self.use_attn_dst)


class QCustomGATConv(CustomGATConv):
    def __init__(self, *args, **kwargs):
        super(QCustomGATConv, self).__init__(*args, **kwargs)
        if isinstance(self.in_channels, int):
            self.lin_l = QLinear(self.in_channels, self.heads * self.out_channels, bias=False)
            self.lin_r = self.lin_l
        else:
            self.lin_l = QLinear(self.in_channels[0], self.heads * self.out_channels, False)
            self.lin_r = QLinear(self.in_channels[1], self.heads * self.out_channels, False)

        if self.residual:
            self.res_fc = Linear(self.in_channels, self.heads * self.out_channels, bias=False)
        else:
            self.register_buffer("res_fc", None)
        self.dropout_module = QDropout(p=self.dropout)

    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
                index: Tensor, ptr: OptTensor,
                size_i: Optional[int]) -> Tensor:
        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
        alpha = ext_quantization.act_quantized_leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = self.dropout_module(alpha)
        return x_j * alpha.unsqueeze(-1)
exact/build/lib.linux-x86_64-cpython-37/exact/gatconv.py
saintslab/i-Exact
from torch._C import qscheme
from torch_geometric.nn.conv import GCNConv, SAGEConv, GCN2Conv, GraphConv, GENConv
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from torch_geometric.typing import Adj, OptTensor, Optional, Union, OptPairTensor, Size
from torch_geometric.utils import softmax
from torch_sparse import SparseTensor

import torch
import torch.nn.functional as F
from torch import Tensor
from torch.nn import Parameter, Sequential
from torch_scatter import scatter, scatter_softmax

from exact.ops import rp2input, input2rp
import exact.cpp_extension.quantization as ext_quantization
from .utils import compute_tensor_bytes, get_memory_usage
from .ops import qlinear, qbatch_norm, qelu, quantize_activation
from .conf import config
from .qsheme import QScheme
from .spmm import qmatmul, QMatmul

QGCNConv_layer_ct = 0
total_act_mem = 0
GPU = 0


def qaddmm(input, mat1, mat2, beta=1., alpha=1):
    return beta * input + alpha * qlinear.apply(mat1, mat2, None)



class QMLP(Sequential):
    def __init__(self, channels, norm,
                 bias=True, dropout=0.):
        m = []
        for i in range(1, len(channels)):
            m.append(QLinear(channels[i - 1], channels[i], bias))

            if i < len(channels) - 1:
                if norm and norm == 'batch':
                    m.append(QBatchNorm1d(channels[i], affine=True))
                elif norm:
                    raise NotImplementedError(
                        f'Normalization layer "{norm}" not supported.')
                m.append(QReLU())
                m.append(QDropout(dropout))
        super(QMLP, self).__init__(*m)


class QGENConv(GENConv):
    def __init__(self, *args, num_layers=2, norm='batch', **kwargs):
        super().__init__(*args, norm=norm, num_layers=num_layers, **kwargs)
        channels = [self.in_channels]
        for _ in range(num_layers - 1):
            channels.append(self.in_channels * 2)
        channels.append(self.out_channels)
        self.mlp = QMLP(channels, norm)
        self.reset_parameters()


    def message(self, x_j: Tensor, edge_attr: OptTensor) -> Tensor:
        msg = x_j if edge_attr is None else x_j + edge_attr
        return ext_quantization.act_quantized_relu(msg) + self.eps


class QGCN2Conv(GCN2Conv):
    def __init__(self, *args, **kwargs):
        super(QGCN2Conv, self).__init__(*args, **kwargs)


    def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""

        if self.normalize:
            if isinstance(edge_index, Tensor):
                cache = self._cached_edge_index
                if cache is None:
                    edge_index, edge_weight = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim), False,
                        self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_edge_index = (edge_index, edge_weight)
                else:
                    edge_index, edge_weight = cache[0], cache[1]

            elif isinstance(edge_index, SparseTensor):
                cache = self._cached_adj_t
                if cache is None:
                    edge_index = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim), False,
                        self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_adj_t = edge_index
                else:
                    edge_index = cache

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        x = self.propagate(edge_index, x=x, edge_weight=edge_weight, size=None)

        x.mul_(1 - self.alpha)
        x_0 = self.alpha * x_0[:x.size(0)]

        if self.weight2 is None:
            out = x.add_(x_0)
            out = qaddmm(out, out, self.weight1, beta=1. - self.beta,
                               alpha=self.beta)
        else:
            out = qaddmm(x, x, self.weight1, beta=1. - self.beta,
                               alpha=self.beta)
            out += qaddmm(x_0, x_0, self.weight2, beta=1. - self.beta,
                                alpha=self.beta)

        return out

    def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
            return edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        return qmatmul(adj_t, x, reduce=self.aggr)


class QSAGEConv(SAGEConv):
    def __init__(self, *args, **kwargs):
        super(QSAGEConv, self).__init__(*args, **kwargs)
        in_channels = self.in_channels
        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)
        # reset the lin_l and lin_l from nn.Linear to Qlinear
        bias = self.lin_l.bias is not None
        self.lin_l = QLinear(in_channels[0], self.out_channels, bias=bias, rp=True)
        if self.root_weight:
            self.lin_r = QLinear(in_channels[1], self.out_channels, bias=False, rp=True)
        self.msg_and_aggr_func = QMatmul(self.aggr)
        if config.single_precision:
            self.scheme = QScheme('sage')
        else:
            self.scheme = None
        self.reset_parameters()


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None) -> Tensor:
        """"""
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
            l_eq_r = True
        else:
            l_eq_r = False
        
        # To avoid wasting memory, 
        # if x is a tensor, we quantize it only once and pass it to the qmatmul and lin_r.
        if l_eq_r:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * x[0].shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(x[0], kept_acts)
            else:
                dim_reduced_input, randmat = x[0], None
            quantized = quantize_activation(dim_reduced_input, self.scheme)
        else:
            quantized, randmat = None, None
        # quantized = None
        # propagate_type: (x: OptPairTensor)
        out = self.propagate(edge_index, x=x, size=size, quantized=quantized, randmat=randmat)
        out = self.lin_l(out)

        x_r = x[1]
        if self.root_weight and x_r is not None:
            out += self.lin_r(x_r, quantized, randmat)

        if self.normalize:
            out = F.normalize(out, p=2., dim=-1)

        return out

    def message_and_aggregate(self, adj_t, x, quantized=None, randmat=None) -> Tensor:
        # adj_t = adj_t.set_value(None, layout=None)
        return self.msg_and_aggr_func(adj_t, x[0], quantized=quantized, randmat=randmat)


class QGraphConv(GraphConv):
    def __init__(self, *args, **kwargs):
        super(QGraphConv, self).__init__(*args, **kwargs)
        in_channels = self.in_channels
        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)
        bias = self.lin_l.bias is not None
        self.lin_l = QLinear(in_channels[0], self.out_channels, bias=bias, rp=True)
        self.lin_r = QLinear(in_channels[1], self.out_channels, bias=False, rp=True)
        self.msg_and_aggr_func = QMatmul(self.aggr)
        self.reset_parameters()


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                edge_weight: OptTensor = None, size: Size = None) -> Tensor:
        """"""
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
        # propagate_type: (x: OptPairTensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=size)
        out = self.lin_l(out)

        x_r = x[1]
        if x_r is not None:
            out += self.lin_r(x_r)
        return out


    def message_and_aggregate(self, adj_t: SparseTensor,
                              x: OptPairTensor) -> Tensor:
        return self.msg_and_aggr_func(adj_t, x[0])


class QGCNConv(GCNConv):
    def __init__(self, *args, **kwargs):
        super(QGCNConv, self).__init__(*args, **kwargs)
        self.weight = Parameter(torch.Tensor(self.out_channels, self.in_channels))
        self.msg_and_aggr_func = QMatmul(self.aggr)
        if config.single_precision:
            self.lin_scheme = QScheme('linear')
        else:
            self.lin_scheme = None
        self.reset_parameters()

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""
        if self.normalize:
            if isinstance(edge_index, Tensor):
                cache = self._cached_edge_index
                if cache is None:
                    edge_index, edge_weight = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim),
                        self.improved, self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_edge_index = (edge_index, edge_weight)
                else:
                    edge_index, edge_weight = cache[0], cache[1]

            elif isinstance(edge_index, SparseTensor):
                cache = self._cached_adj_t
                if cache is None:
                    edge_index = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim),
                        self.improved, self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_adj_t = edge_index
                else:
                    edge_index = cache
        if self.training:
            x = qlinear.apply(x, self.weight, None, None, self.lin_scheme)
        else:
            x = F.linear(x, self.weight, None)

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=None)

        if self.bias is not None:
            out += self.bias
        return out

    def message_and_aggregate(self, adj_t, x):
        return self.msg_and_aggr_func(adj_t, x)


class QLinear(torch.nn.Linear):
    def __init__(self, input_features, output_features, bias=True, rp=True):
        super(QLinear, self).__init__(input_features, output_features, bias)
        if config.single_precision:
            self.scheme = QScheme('linear')
        else:
            self.scheme = None
        self.rp = rp

    def forward(self, input, quantized=None, randmat=None):
        if self.training:
            return qlinear.apply(input, self.weight, self.bias, quantized, randmat, self.scheme, self.rp)
        else:
            return super(QLinear, self).forward(input)


class QReLU(torch.nn.Module):
    def __init__(self, inplace=False):
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_relu(input)
        else:
            return F.relu(input)


class QLeakyReLU(torch.nn.Module):
    def __init__(self, negative_slope, inplace=False):
        super().__init__()
        self.negative_slope = negative_slope

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return ext_quantization.act_quantized_leaky_relu(input, self.negative_slope)


class QELU(torch.nn.Module):
    def __init__(self, alpha, inplace=False):
        super().__init__()
        self.alpha = alpha

    def forward(self, input):
        if self.training:
            return qelu.apply(input, self.alpha)
        else:
            return F.elu(input, self.alpha)


class QDropout(torch.nn.Module):
    def __init__(self, p):
        self.p = p
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_dropout(input, self.p, self.training)
        else:
            return F.dropout(input, self.p, self.training)


class QDropout2(torch.nn.Module):
    def __init__(self, p):
        self.p = p
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_dropout(input, self.p, self.training) * (1 - self.p)
        else:
            return F.dropout(input, self.p, self.training) * (1 - self.p)


class QBatchNorm1d(torch.nn.BatchNorm1d):
    def __init__(self, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
        super(QBatchNorm1d, self).__init__(num_features, eps, momentum, affine, track_running_stats)
        if config.single_precision:
            self.scheme = QScheme('batchnorm')
        else:
            self.scheme = None

    def forward(self, input):
        if not self.training:
            return super(QBatchNorm1d, self).forward(input)

        self._check_input_dim(input)

        # exponential_average_factor is set to self.momentum
        # (when it is available) only so that it gets updated
        # in ONNX graph when this node is exported to ONNX.
        if self.momentum is None:
            exponential_average_factor = 0.0
        else:
            exponential_average_factor = self.momentum

        if self.training and self.track_running_stats:
            # TODO: if statement only here to tell the jit to skip emitting this when it is None
            if self.num_batches_tracked is not None:
                self.num_batches_tracked = self.num_batches_tracked + 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        """ Decide whether the mini-batch stats should be used for normalization rather than the buffers.
        Mini-batch stats are used in training mode, and in eval mode when buffers are None.
        """
        if self.training:
            bn_training = True
        else:
            bn_training = (self.running_mean is None) and (self.running_var is None)

        """Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
        used for normalization (i.e. in eval mode when buffers are not None).
        """
        return qbatch_norm.apply(
            input,
            # If buffers are not to be tracked, ensure that they won't be updated
            self.running_mean if not self.training or self.track_running_stats else None,
            self.running_var if not self.training or self.track_running_stats else None,
            self.weight, self.bias, bn_training, exponential_average_factor, self.eps, self.scheme)

exact/build/lib.linux-x86_64-cpython-37/exact/layers.py
saintslab/i-Exact
from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict
from collections import OrderedDict
from torch_geometric.nn.conv.gen_conv import GENConv
from torch_geometric.nn.conv.graph_conv import GraphConv
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor, device, dtype

from torch_geometric.nn.conv import GCNConv, SAGEConv, GCN2Conv, SGConv
from .layers import QBatchNorm1d, QLinear, QReLU, QGCNConv, QDropout, QDropout2, QSAGEConv, QGENConv, QGraphConv, QGCN2Conv
from .conf import config
from .gatconv import CustomGATConv, QCustomGATConv


class QModule(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        QModule.convert_layers(model)

    @staticmethod
    def convert_layers(module):
        for name, child in module.named_children():
            # Do not convert layers that are already quantized
            if isinstance(child, (QBatchNorm1d, QLinear, QReLU, QGCNConv, QGCN2Conv, QDropout, QSAGEConv, QGENConv, QCustomGATConv)):
                continue
            if isinstance(child, nn.BatchNorm1d) and config.enable_quantized_bn:
                setattr(module, name, QBatchNorm1d(child.num_features, child.eps, child.momentum,
                    child.affine, child.track_running_stats))
            elif isinstance(child, nn.Linear):
                setattr(module, name, QLinear(child.in_features, child.out_features,
                    child.bias is not None))
            elif isinstance(child, nn.ReLU):
                setattr(module, name, QReLU())
            elif isinstance(child, nn.Dropout):
                if config.dropout2:
                    setattr(module, name, QDropout2(child.p))
                else:
                    setattr(module, name, QDropout(child.p))
            elif isinstance(child, GCNConv):
                setattr(module, name, QGCNConv(child.in_channels, child.out_channels, child.improved, child.cached,
                                               child.add_self_loops, child.normalize, child.bias is not None,
                                               aggr=child.aggr))
            elif isinstance(child, GCN2Conv):
                beta = child.beta
                shared_weights = child.weight2 is None
                setattr(module, name, QGCN2Conv(child.channels, alpha=child.alpha, theta=None, layer=None, shared_weights=shared_weights,
                                                cached=child.cached, add_self_loops=child.add_self_loops, normalize=child.normalize))
                curconv = getattr(module, name)
                curconv.beta = child.beta
            elif isinstance(child, SAGEConv):
                setattr(module, name, QSAGEConv(child.in_channels, child.out_channels, child.normalize, child.root_weight,
                                               child.lin_l.bias is not None))
            elif isinstance(child, GraphConv):
                setattr(module, name, QGraphConv(child.in_channels, child.out_channels, child.aggr, child.lin_l.bias is not None))
            elif isinstance(child, CustomGATConv):
                setattr(module, name, QCustomGATConv(child.in_channels, child.out_channels, child.heads, child.concat, 
                                                     child.negative_slope, child.dropout, child.add_self_loops, child.bias is not None, 
                                                     child.residual, child.use_attn_dst))
            elif isinstance(child, GENConv):
                msg_norm = child.msg_norm is not None
                learn_msg_scale = True if (msg_norm and child.msg_norm.scale.requires_grad) else False
                learn_p = isinstance(child.p, torch.nn.Parameter)
                is_softmax = child.aggr == 'softmax'
                if is_softmax and isinstance(child.t, torch.nn.Parameter):
                    learn_t = True
                else:
                    learn_t = False
                num_layers = 0
                norm = 'batch'
                for m in child.mlp:
                    if isinstance(m, torch.nn.Linear):
                        num_layers += 1
                    if isinstance(m, (nn.BatchNorm1d, nn.LayerNorm, nn.InstanceNorm1d)):
                        if isinstance(m, nn.BatchNorm1d):
                            pass
                        elif isinstance(m, nn.LayerNorm):
                            norm = 'layer'
                        elif isinstance(m, nn.InstanceNorm1d):
                            norm = 'instance'

                setattr(module, name, QGENConv(child.in_channels, child.out_channels, child.aggr, child.initial_t, learn_t, child.initial_p, learn_p, 
                                               msg_norm, learn_msg_scale, norm, num_layers, child.eps))
            else:
                QModule.convert_layers(child)

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor]],
                        strict: bool = True):
        # remove the prefix "model." added by this wrapper
        new_state_dict = OrderedDict([("model." + k,  v) for k, v in state_dict.items()])
        return super().load_state_dict(new_state_dict, strict)

    def state_dict(self, destination=None, prefix='', keep_vars=False):
        ret = super().state_dict(destination, prefix, keep_vars)

        # remove the prefix "model." added by this wrapper
        ret = OrderedDict([(k[6:], v) for k, v in ret.items()])
        return ret

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


    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        return self.model.mini_inference(x_all, loader)
exact/build/lib.linux-x86_64-cpython-37/exact/module.py
saintslab/i-Exact
import pdb
import numpy as np
import os
import csv

import torch
import torch.nn.functional as F
from bitsandbytes import functional as bnb_F
from torch.autograd.function import Function
from torch.cuda.amp import custom_fwd, custom_bwd

from torch_sparse import matmul
from .conf import config
import exact.cpp_extension.spmm as spmm 

import exact.cpp_extension.backward_func as ext_backward_func
import exact.cpp_extension.quantization as ext_quantization

from exact.utils import empty_cache, get_memory_usage, compute_tensor_bytes, swap_to_cpu, cast_low_bit_int


def quantize_and_pack(data, bits, mn, mx):
    if config.simulate:
        N = data.shape[0]
        output = data  
        B = 2 ** bits - 1
        mn = mn - 1e-6
        mx = mx + 1e-6
        scale = B / (mx - mn) 
        output = (output - mn.view(-1, 1)) * scale.view(-1, 1)
        if config.stochastic:
            noise = output.new(output.shape).uniform_(-0.5, 0.5)
            output.add_(noise)

        output = F.relu(output)
        output = output.round_().int()
    else:
        # Pack to bitstream
        assert type(bits) == int
        pack_func = ext_quantization.pack_single_precision
        scale = (2 ** bits - 1) / (mx - mn) 
        if config.use_optimal_lo:
            dim = data.size(-1)
            a, b, c, d, f = 0.83162565, 0.02993653, 116.05528511, 0.00001225, -0.00171617
            config.lo = a + b * np.log(dim * c) + d * dim + f * dim**0.5
        lo, hi = config.lo, 3 - config.lo
        output = pack_func(data, mn, mx, scale.to(data.dtype), bits, config.stochastic, lo, hi)
        if config.swap:
            output = swap_to_cpu(output)
    return output, scale, (lo, hi)


def dequantize_and_unpack(data, shape, bits, scale, mn, lo, hi):
    if config.simulate:
        data = data / scale.view(-1, 1) + mn.view(-1, 1)
    else:
        if config.swap:
            data = data.cuda()
        N = shape[0]
        num_features = int(np.prod(shape[1:]))
        # Unpack bitstream
        assert type(bits) == int
        unpack_func = ext_quantization.unpack_single_precision
        data = unpack_func(data, bits, scale, mn, N, num_features, lo, hi)
    return data


def no_scheme_compute_quantization_bits(input):
    N = input.shape[0]
    input_flatten = input.view(N, -1)

    mn, mx = torch.min(input_flatten, 1)[0], torch.max(input_flatten, 1)[0]

    b = config.activation_compression_bits[0]
    return input_flatten, b, mn, mx


def closest_multiple(col_size, multiple):
    remainder = col_size % multiple
    if remainder > multiple / 2:
        return col_size + (multiple - remainder)
    else:
        return col_size - remainder


def quantize_activation(input, scheme):
    #input = expand_middle(input)
    # col_size is actually col_size multiplier
    col_size = config.col_size * input.size(-1)
    #col_size = max(col_size, input.size(-1))
    col_size = np.clip(col_size, input.size(-1), 1024)
    col_size = closest_multiple(col_size, input.size(-1))
    
    # Calculate the size of the reshaped tensor
    elements = input.numel()
    remaining_elements = elements % col_size
    main_part_rows = (elements - remaining_elements) // input.size(1)

    orig_feat_dim = input.size(-1)
    input2 = input[main_part_rows:].clone()

    #input = input.narrow(0, 0, main_part_rows).view(-1, col_size)
    input = input[:main_part_rows].view(-1, col_size)

    if not config.compress_activation:
        if config.swap:
            input = swap_to_cpu(input)
        return input, None, None, None

    if scheme:
        input_groups, q_bits, q_min, mx = scheme.compute_quantization_bits(input)
    else:
        input_groups, q_bits, q_min, mx = no_scheme_compute_quantization_bits(input)

    q_input, q_scale, (lo, hi) = quantize_and_pack(input_groups, q_bits, q_min, mx)
    if input.dtype == torch.float32:
        return q_input, q_bits, q_scale.to(torch.bfloat16), q_min.to(torch.bfloat16), input2, input.shape, orig_feat_dim, (lo, hi)
    else:
        return q_input, q_bits, q_scale, q_min, input2, input.shape, orig_feat_dim, (lo, hi)


def dequantize_activation(quantized, q_input_shape):
    if not config.compress_activation:
        ret = quantized[0]
        if config.swap:
            ret = ret.cuda(non_blocking=True)
        return ret

    q_input, q_bits, q_scale, q_min, input2, q_input_shape, orig_feat_dim, (lo, hi) = quantized
    if q_scale.dtype == torch.bfloat16:
        q_scale = q_scale.to(torch.float32)
        q_min = q_min.to(torch.float32)
    input = dequantize_and_unpack(q_input, q_input_shape, q_bits, q_scale, q_min, lo, hi)     

    input = input.view(-1, orig_feat_dim)
    input = torch.cat([input, input2], dim=0)
    return input.contiguous()


linear_layer_ct = 0
qmatmul_layer_ct = 0
bn_layer_ct = 0
total_act_mem = 0
GPU = 0

@torch.no_grad()
def input2rp(input, kept_acts):
    assert len(input.size()) == 2
    rand_mat_size = (input.shape[1], kept_acts)
    # Create random matrix
    def gen_rad_mat(rm_size, feat_size, device, dtype):
        bern = torch.randint(2, size=rm_size, device=device, requires_grad=False, dtype=dtype)
        return (2.0 * bern - 1) / feat_size **0.5

    rand_matrix = gen_rad_mat(rand_mat_size, kept_acts, input.device, input.dtype)
    dim_reduced_input = torch.matmul(input, rand_matrix)
    return dim_reduced_input, rand_matrix


@torch.no_grad()
def rp2input(dim_reduced_input, input_shape, rand_matrix):
    assert len(dim_reduced_input.size()) == 2
    input = torch.matmul(dim_reduced_input, rand_matrix.t())    
    return input.view(input_shape)


class qlinear(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, input, weight, bias=None, quantized=None, randmat=None, scheme=None, rp=True):
        if quantized is not None:
            if randmat is None:
                assert rp is False or config.kept_frac == 1.0
                # quantized somewhere before
                ori_input_shape, proj_input_shape = input.shape, input.shape
            else:
                assert (rp is True and config.kept_frac < 1.0) and (input.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = input.shape, torch.Size([input.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0 and rp:
                kept_acts = int(config.kept_frac * input.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(input, kept_acts)
                ori_input_shape, proj_input_shape = input.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = input, None
                ori_input_shape, proj_input_shape = input.shape, input.shape
            quantized = quantize_activation(dim_reduced_input, scheme)

        empty_cache(config.empty_cache_threshold)

        ctx.scheme = scheme
        ctx.saved = quantized, weight, bias, randmat
        ctx.other_args = ori_input_shape, proj_input_shape, rp
        res = F.linear(input, weight, bias)
        return res

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        quantized, weight, bias, randmat = ctx.saved
        ori_input_shape, q_input_shape, rp = ctx.other_args
        input = dequantize_activation(quantized, q_input_shape)
        if config.kept_frac < 1.0 and rp:
            input = rp2input(input, ori_input_shape, randmat)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_input, grad_weight, grad_bias = ext_backward_func.linear_backward(grad_output, input, weight, bias)
        # grad_input = grad_output.mm(weight)
        # grad_weight = grad_output.t().mm(input)
        # if bias is not None:
        #     grad_bias = grad_output.sum(0)
        # else:
        #     grad_bias = None
        del input, grad_output
        empty_cache(config.empty_cache_threshold)
        return grad_input, grad_weight, grad_bias, None, None, None, None



class qelu(Function):
    @staticmethod
    def forward(ctx, input, alpha, scheme=None):
        quantized = quantize_activation(input, scheme)
        empty_cache(config.empty_cache_threshold)
        ctx.scheme = scheme
        ctx.saved = quantized
        ctx.other_args = input.shape, alpha
        res = F.elu(input, alpha)
        return res

    @staticmethod
    def backward(ctx, grad_output):
        # if ctx.scheme:
        #     ctx.scheme.set_scale(grad_output)

        quantized = ctx.saved
        q_input_shape, alpha = ctx.other_args
        input = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_input = ext_backward_func._elu_backward_cuda(grad_output, input, alpha)
        return grad_input, None, None

class qbatch_norm(Function):
    @staticmethod
    def forward(ctx, input, running_mean, running_var, weight, bias,
                training, exponential_average_factor, eps, scheme):
        quantized = quantize_activation(input, scheme)

        if training:
            output, save_mean, save_var, reserve, _ = ext_backward_func._batch_norm_impl_index(input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps, True)

        else:
            output, save_mean, save_var = ext_backward_func.native_batch_norm(
                input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps)
            reserve = None

        # output, save_mean, save_var = ext_backward_func.native_batch_norm(
        #     input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps)
        # reserve = None
        ctx.scheme = scheme
        ctx.other_args = input.shape
        ctx.saved = (quantized, weight, running_mean, running_var, save_mean, save_var, training, eps, reserve)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        quantized, weight, running_mean, running_var, save_mean, save_var, training, eps, reserve = ctx.saved

        q_input_shape = ctx.other_args

        input = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved

        empty_cache(config.empty_cache_threshold)

        # if training:
        #     input = input.contiguous()
        #     grad_input, grad_weight, grad_bias = ext_backward_func.cudnn_batch_norm_backward(
        #         input, grad_output, weight, running_mean, running_var, save_mean, save_var, eps, reserve)
        # else:
        #     grad_input, grad_weight, grad_bias = ext_backward_func.native_batch_norm_backward(
        #         grad_output, input, weight, running_mean, running_var, save_mean, save_var, training, eps,
        #         [ctx.needs_input_grad[0], ctx.needs_input_grad[3], ctx.needs_input_grad[4]]
        #     )

        grad_input, grad_weight, grad_bias = ext_backward_func.native_batch_norm_backward(
            grad_output, input, weight, running_mean, running_var, save_mean, save_var, training, eps,
            [ctx.needs_input_grad[0], ctx.needs_input_grad[3], ctx.needs_input_grad[4]]
        )
        return grad_input, None, None, grad_weight, grad_bias, None, None, None, None


class qspmm_sum(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, row, rowptr, col, value, colptr, csr2csc, has_value, other, quantized=None, randmat=None, scheme=None):
        result = spmm.spmm_sum_fw(row, rowptr, col, value, colptr, csr2csc, other)
        if quantized is not None:
            if randmat is None:
                assert config.kept_frac == 1.0
                # quantized somewhere before
                ori_input_shape, proj_input_shape = other.shape, other.shape
            else:
                assert config.kept_frac < 1.0 and (other.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = other.shape, torch.Size([other.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * other.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(other, kept_acts)
                ori_input_shape, proj_input_shape = other.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = other, None
                ori_input_shape, proj_input_shape = other.shape, other.shape
            quantized = quantize_activation(dim_reduced_input, scheme)
        
        empty_cache(config.empty_cache_threshold)
        ctx.saved = row, rowptr, col, value, colptr, csr2csc, quantized, randmat
        ctx.other_args = has_value, ori_input_shape, proj_input_shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.scheme = scheme
        return result


    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        row, rowptr, col, value, colptr, csr2csc, quantized, randmat = ctx.saved
        row = col if row is None else row
        value = col if value is None else value
        colptr = col if colptr is None else colptr
        csr2csc = col if csr2csc is None else csr2csc
        has_value, ori_input_shape, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        if config.kept_frac < 1.0:
            other = rp2input(other, ori_input_shape, randmat)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_sum_bw(row, rowptr, col, value, colptr, csr2csc, other, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        del other
        empty_cache(config.empty_cache_threshold)
        return None, None, None, grad_value, None, None, None, grad_mat, None, None, None


class qspmm_mean(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        result = spmm.spmm_mean_fw(row, rowptr, col, value, rowcount, colptr, csr2csc, other)
        if quantized is not None:
            if randmat is None:
                assert config.kept_frac == 1.0
                # quantized somewhere before without random projection
                ori_input_shape, proj_input_shape = other.shape, other.shape
            else:
                assert config.kept_frac < 1.0 and (other.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = other.shape, torch.Size([other.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * other.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(other, kept_acts)
                ori_input_shape, proj_input_shape = other.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = other, None
                ori_input_shape, proj_input_shape = other.shape, other.shape
            quantized = quantize_activation(dim_reduced_input, scheme)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = row, rowptr, col, value, rowcount, colptr, csr2csc, quantized, randmat
        ctx.other_args = has_value, ori_input_shape, proj_input_shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.scheme = scheme
        return result

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        row, rowptr, col, value, rowcount, colptr, csr2csc, quantized, randmat = ctx.saved
        row = col if row is None else row
        value = col if value is None else value
        rowcount = col if rowcount is None else rowcount
        colptr = col if colptr is None else colptr
        csr2csc = col if csr2csc is None else csr2csc
        has_value, ori_input_shape, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        # here is one ugly trick: if we know value does not need gradient,
        # we actually do not need the ``other'' matrix to calculate the gradient.
        # So here we just pass a dummy matrix to the CUDA kernel.
        # TODO: engineering optimization.
        if value_requires_grad:
            other = dequantize_activation(quantized, q_input_shape)
            if config.kept_frac < 1.0:
                other = rp2input(other, ori_input_shape, randmat)
            # other = rp2input(other, ori_input_shape, randmat)

        else:
            if quantized[2] is None:
                other = torch.tensor([1.], dtype=torch.float, device='cuda')
            else:
                if quantized[2].dtype == torch.bfloat16:
                    dtype = torch.float
                else:
                    dtype = quantized[2].dtype 
                other = torch.tensor([1.], dtype=dtype, device=quantized[2].device)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_mean_bw(row, rowptr, col, value, rowcount, colptr, csr2csc, other, grad_outputs, 
                                                 has_value, value_requires_grad, mat_requires_grad)
        del other
        empty_cache(config.empty_cache_threshold)
        return None, None, None, grad_value, None, None, None, None, grad_mat, None, None, None


class qspmm_max(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, rowptr, col, value, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        output, arg_out = spmm.spmm_max_fw(rowptr, col, value, other)
        if quantized is None:
            quantized = quantize_activation(other, scheme)
        else:
            assert isinstance(quantized, tuple)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = col, value, quantized, arg_out
        ctx.other_args = has_value, other.shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.mark_non_differentiable(arg_out)
        ctx.scheme = scheme
        return output

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        col, value, quantized, arg_out = ctx.saved
        value = col if value is None else value
        has_value, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_max_bw(col, value, other, arg_out, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        return None, None, grad_value, None, grad_mat, None, None, None


class qspmm_min(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, rowptr, col, value, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        output, arg_out =  spmm.spmm_min_fw(rowptr, col, value, other)
        if quantized is None:
            quantized = quantize_activation(other, scheme)
        else:
            assert isinstance(quantized, tuple)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = col, value, quantized, arg_out
        ctx.other_args = has_value, other.shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.mark_non_differentiable(arg_out)
        ctx.scheme = scheme
        return output
   
    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        col, value, quantized, arg_out = ctx.saved
        value = col if value is None else value
        has_value, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        # if ctx.scheme:
        #     ctx.scheme.set_scale(grad_outputs)
        grad_value, grad_mat = spmm.spmm_min_bw(col, value, other, arg_out, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        return None, None, grad_value, None, grad_mat, None, None, None
exact/build/lib.linux-x86_64-cpython-37/exact/ops.py
saintslab/i-Exact
import pdb

import torch
from .conf import config
import exact.cpp_extension.calc_precision as ext_calc_precision


class QScheme(object):
    num_samples = 1
    num_matmul_layer = 0
    num_lin_layer = 0
    num_batchnorm_layer = 0
    num_sage_layer = 0
    batch = None
    update_scale = True
    layers = []
    prev_layer = None

    def __init__(self, name):
        self.initial_bits = config.initial_bits
        self.bits = config.activation_compression_bits[0]
        QScheme.layers.append(self)

        assert name in ['linear', 'matmul', 'batchnorm', 'sage']
        if name == 'linear':
            self.name = f'linear_{QScheme.num_lin_layer}'
            QScheme.num_lin_layer += 1
        elif name == 'matmul':
            self.name = f'matmul_{QScheme.num_matmul_layer}'
            QScheme.num_matmul_layer += 1
        elif name == 'batchnorm':
            self.name = f'batchnorm_{QScheme.num_batchnorm_layer}'
            QScheme.num_batchnorm_layer += 1
        elif name == 'sage':
            self.name = f'sage_{QScheme.num_sage_layer}'
            QScheme.num_sage_layer += 1
        else:
            raise ValueError


    def compute_quantization_bits(self, input):
        QScheme.prev_layer = self
        N = input.shape[0]
        input_flatten = input.view(N, -1)
        mn, mx = torch.min(input_flatten, 1)[0], torch.max(input_flatten, 1)[0]
        # range = mx - mn
        b = config.activation_compression_bits[0]
        return input_flatten, b, mn, mx
exact/build/lib.linux-x86_64-cpython-37/exact/qsheme.py
saintslab/i-Exact
import torch
from .conf import config
from .qsheme import QScheme
from .ops import qspmm_sum, qspmm_mean, qspmm_max, qspmm_min


class QMatmul(torch.nn.Module):
    def __init__(self, reduce):
        super().__init__()
        if reduce in ['sum', 'add']:
            reduce = 'sum'
        self.reduce = reduce
        if config.single_precision:
            self.scheme = QScheme('matmul')
        else:
            self.scheme = None

    def forward(self, src, other, quantized=None, randmat=None): 
        rowptr, col, value = src.csr()
        has_value = value is not None
        if self.reduce == 'sum':
            row = src.storage._row
            csr2csc = src.storage._csr2csc
            colptr = src.storage._colptr
            if has_value and value.requires_grad:
                row = src.storage.row()
            if other.requires_grad:
                row = src.storage.row()
                csr2csc = src.storage.csr2csc()
                colptr = src.storage.colptr()
            return qspmm_sum.apply(row, rowptr, col, value, colptr, csr2csc, has_value, other, 
                                quantized, randmat, self.scheme)
        elif self.reduce == 'mean':
            row = src.storage._row
            rowcount = src.storage._rowcount
            csr2csc = src.storage._csr2csc
            colptr = src.storage._colptr
            if has_value and value.requires_grad:
                row = src.storage.row()
            if other.requires_grad:
                row = src.storage.row()
                rowcount = src.storage.rowcount()
                csr2csc = src.storage.csr2csc()
                colptr = src.storage.colptr()
            return qspmm_mean.apply(row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                                    quantized, randmat, self.scheme)
        elif self.reduce == 'min':
            return qspmm_min.apply(rowptr, col, value, has_value, other, quantized, randmat, self.scheme)
        elif self.reduce == 'max':
            return qspmm_max.apply(rowptr, col, value, has_value, other, quantized, randmat, self.scheme)
        else:
            raise ValueError


def qmatmul(src, other, reduce="sum", quantized=None, randmat=None, scheme=None):
    rowptr, col, value = src.csr()
    has_value = value is not None
    if reduce in ['sum', 'add']:
        reduce = 'sum'
        row = src.storage._row
        csr2csc = src.storage._csr2csc
        colptr = src.storage._colptr
        if has_value and value.requires_grad:
            row = src.storage.row()
        if other.requires_grad:
            row = src.storage.row()
            csr2csc = src.storage.csr2csc()
            colptr = src.storage.colptr()
        return qspmm_sum.apply(row, rowptr, col, value, colptr, csr2csc, has_value, other, 
                               quantized, randmat, scheme)
    elif reduce == 'mean':
        row = src.storage._row
        rowcount = src.storage._rowcount
        csr2csc = src.storage._csr2csc
        colptr = src.storage._colptr
        if has_value and value.requires_grad:
            row = src.storage.row()
        if other.requires_grad:
            row = src.storage.row()
            rowcount = src.storage.rowcount()
            csr2csc = src.storage.csr2csc()
            colptr = src.storage.colptr()
        return qspmm_mean.apply(row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                                quantized, randmat, scheme)
    elif reduce == 'min':
        return qspmm_min.apply(rowptr, col, value, has_value, other, quantized, randmat, scheme)
    elif reduce == 'max':
        return qspmm_max.apply(rowptr, col, value, has_value, other, quantized, randmat, scheme)
    else:
        raise ValueError
exact/build/lib.linux-x86_64-cpython-37/exact/spmm.py
saintslab/i-Exact
import os
from collections import OrderedDict
import json

import torch
import numpy as np


def cast_low_bit_int(idx_vec):
    if idx_vec is None:
        return idx_vec
    assert idx_vec.dtype == torch.int64
    max_v = torch.max(idx_vec)
    if max_v <= np.iinfo(np.int16).max:
        return idx_vec.type(torch.int16)
    elif max_v < np.iinfo(np.int32).max:
        return idx_vec.type(torch.int32)
    else:
        return idx_vec


def cast_adj(data):
    data.adj_t.storage._col = cast_low_bit_int(data.adj_t.storage._col)
    data.adj_t.storage._row = cast_low_bit_int(data.adj_t.storage._row)
    data.adj_t.storage._rowptr = cast_low_bit_int(data.adj_t.storage._rowptr)
    data.adj_t.storage._rowcount = cast_low_bit_int(data.adj_t.storage._rowcount)
    data.adj_t.storage._colptr = cast_low_bit_int(data.adj_t.storage._colptr)
    data.adj_t.storage._csr2csc = cast_low_bit_int(data.adj_t.storage._csr2csc)


def swap_to_cpu(tensor):
    tensor_cpu = torch.empty(tensor.shape, dtype=tensor.dtype, device='cpu', pin_memory=True)
    tensor_cpu.copy_(tensor, non_blocking=True)
    return tensor_cpu


def get_memory_usage(gpu, print_info=False):
    """Get accurate gpu memory usage by querying torch runtime"""
    allocated = torch.cuda.memory_allocated(gpu)
    reserved = torch.cuda.memory_reserved(gpu)
    if print_info:
        print("allocated: %.2f MB" % (allocated / 1024 / 1024), flush=True)
        print("reserved:  %.2f MB" % (reserved / 1024 / 1024), flush=True)
    return allocated


def compute_tensor_bytes(tensors):
    """Compute the bytes used by a list of tensors"""
    if not isinstance(tensors, (list, tuple)):
        tensors = [tensors]

    ret = 0
    for x in tensors:
        if x.dtype in [torch.int64, torch.long]:
            ret += np.prod(x.size()) * 8
        if x.dtype in [torch.float32, torch.int, torch.int32]:
            ret += np.prod(x.size()) * 4
        elif x.dtype in [torch.bfloat16, torch.float16, torch.int16]:
            ret += np.prod(x.size()) * 2
        elif x.dtype in [torch.int8]:
            ret += np.prod(x.size())
        else:
            print(x.dtype)
            raise ValueError()
    return ret


def empty_cache(ratio):
    if ratio is None:
        return
    allocated = torch.cuda.memory_allocated(0)
    reserved = torch.cuda.memory_reserved(0)
    if reserved > 0 and allocated / reserved < ratio:
        torch.cuda.empty_cache()


def disable_cache_allocator():
    os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'


def enable_cache_allocator():
    del os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING']


class GlobalExpRecorder:
    def __init__(self):
        self.val_dict = OrderedDict()

    def record(self, key, value, float_round=6):
        if isinstance(value, (np.int32, np.int64)):
            value = int(value)
        if isinstance(value, (float, np.float32, np.float64)):
            value = round(value, float_round)

        self.val_dict[key] = value

    def dump(self, filename):
        with open(filename, "a") as fout:
            fout.write(json.dumps(self.val_dict) + '\n')
        print("Save exp results to %s" % filename)

    def clear(self):
        pass

exp_recorder = GlobalExpRecorder()

exact/build/lib.linux-x86_64-cpython-37/exact/utils.py
saintslab/i-Exact
from .utils import get_memory_usage, compute_tensor_bytes, exp_recorder, cast_adj, cast_low_bit_int
from .layers import QGCNConv, QLinear, QReLU, QBatchNorm1d, QDropout, QSAGEConv, QELU, QGCN2Conv, QGENConv
from .conf import config
from .module import QModule
from .gatconv import CustomGATConv, QCustomGATConv
exact/build/lib.linux-x86_64-cpython-38/exact/__init__.py
saintslab/i-Exact
class QuantizationConfig:
    def __init__(self):
        self.compress_activation = True
        self.activation_compression_bits = [2]
        self.initial_bits = 8
        self.stochastic = True
        self.use_gradient = False
        self.simulate = False
        self.enable_quantized_bn = True
        self.kept_frac = 1.0
        self.single_precision = True
        self.dropout2 = False
        self.col_size = 0
        self.last_epoch = False
        self.lo = 1
        self.use_optimal_lo = False
        # Memory management flag
        self.empty_cache_threshold = None
        self.swap = False
        self.amp = False
        # Debug related flag
        self.debug_memory_op_forward = False
        self.debug_memory_op_backward = False
        self.debug_remove_bn = False
        self.debug_remove_relu = False
        self.debug_remove_dropout = False

config = QuantizationConfig()
exact/build/lib.linux-x86_64-cpython-38/exact/conf.py
saintslab/i-Exact
from exact.layers import QDropout, QLinear
from typing import Union, Tuple, Optional
from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
                                    OptTensor)

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import Parameter, Linear
from torch_sparse import SparseTensor, set_diag
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops, softmax

from torch_geometric.nn.inits import zeros
import exact.cpp_extension.quantization as ext_quantization


class CustomGATConv(MessagePassing):
    def __init__(self, in_channels: Union[int, Tuple[int, int]],
                 out_channels: int, heads: int = 1, concat: bool = True,
                 negative_slope: float = 0.2, dropout: float = 0.0,
                 add_self_loops: bool = True, bias: bool = True, 
                 residual=True, use_attn_dst=True, **kwargs):
        kwargs.setdefault('aggr', 'add')
        super(CustomGATConv, self).__init__(node_dim=0, **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.negative_slope = negative_slope
        self.dropout = dropout
        self.add_self_loops = add_self_loops
        self.residual = residual
        self.use_attn_dst = use_attn_dst

        if isinstance(in_channels, int):
            self.lin_l = Linear(in_channels, heads * out_channels, bias=False)
            self.lin_r = self.lin_l
        else:
            self.lin_l = Linear(in_channels[0], heads * out_channels, False)
            self.lin_r = Linear(in_channels[1], heads * out_channels, False)

        if residual:
            self.res_fc = Linear(in_channels, heads * out_channels, bias=False)
        else:
            self.register_buffer("res_fc", None)

        self.att_l = Parameter(torch.Tensor(1, heads, out_channels))
        if self.use_attn_dst:
            self.att_r = Parameter(torch.Tensor(1, heads, out_channels))
        else:
            self.att_r = self.att_l
        if bias and concat:
            self.bias = Parameter(torch.Tensor(heads * out_channels))
        elif bias and not concat:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()


    def reset_parameters(self):
        gain = torch.nn.init.calculate_gain("relu")
        torch.nn.init.xavier_normal_(self.lin_l.weight, gain=gain)
        torch.nn.init.xavier_normal_(self.lin_r.weight, gain=gain)
        if isinstance(self.res_fc, Linear):
            torch.nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
        torch.nn.init.xavier_normal_(self.att_l, gain=gain)
        if self.use_attn_dst:
            torch.nn.init.xavier_normal_(self.att_r, gain=gain)
        zeros(self.bias)


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None):
        H, C = self.heads, self.out_channels

        x_l: OptTensor = None
        x_r: OptTensor = None
        alpha_l: OptTensor = None
        alpha_r: OptTensor = None
        if isinstance(x, Tensor):
            assert x.dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = x_r = self.lin_l(x).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if self.use_attn_dst:
                alpha_r = (x_r * self.att_r).sum(dim=-1)
            else:
                alpha_r = alpha_l
        else:
            x_l, x_r = x[0], x[1]
            assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = self.lin_l(x_l).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if x_r is not None:
                x_r = self.lin_r(x_r).view(-1, H, C)
                alpha_r = (x_r * self.att_r).sum(dim=-1)

        assert x_l is not None
        assert alpha_l is not None

        if self.add_self_loops:
            if isinstance(edge_index, Tensor):
                num_nodes = x_l.size(0)
                if x_r is not None:
                    num_nodes = min(num_nodes, x_r.size(0))
                if size is not None:
                    num_nodes = min(size[0], size[1])
                edge_index, _ = remove_self_loops(edge_index)
                edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
            elif isinstance(edge_index, SparseTensor):
                edge_index = set_diag(edge_index)

        # propagate_type: (x: OptPairTensor, alpha: OptPairTensor)
        out = self.propagate(edge_index, x=(x_l, x_r),
                             alpha=(alpha_l, alpha_r), size=size)


        if self.concat:
            out = out.view(-1, self.heads * self.out_channels)
        else:
            out = out.mean(dim=1)
        if self.residual:
            out += self.res_fc(x)
        if self.bias is not None:
            out += self.bias
        return out

    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
                index: Tensor, ptr: OptTensor,
                size_i: Optional[int]) -> Tensor:
        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
        alpha = F.leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
        return x_j * alpha.unsqueeze(-1)

    def __repr__(self):
        return '{}({}, {}, heads={}, residual={}, use_attn_dst={})'.format(self.__class__.__name__,
                                                                           self.in_channels, self.out_channels, 
                                                                           self.heads, self.residual, 
                                                                           self.use_attn_dst)


class QCustomGATConv(CustomGATConv):
    def __init__(self, *args, **kwargs):
        super(QCustomGATConv, self).__init__(*args, **kwargs)
        if isinstance(self.in_channels, int):
            self.lin_l = QLinear(self.in_channels, self.heads * self.out_channels, bias=False)
            self.lin_r = self.lin_l
        else:
            self.lin_l = QLinear(self.in_channels[0], self.heads * self.out_channels, False)
            self.lin_r = QLinear(self.in_channels[1], self.heads * self.out_channels, False)

        if self.residual:
            self.res_fc = Linear(self.in_channels, self.heads * self.out_channels, bias=False)
        else:
            self.register_buffer("res_fc", None)
        self.dropout_module = QDropout(p=self.dropout)

    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
                index: Tensor, ptr: OptTensor,
                size_i: Optional[int]) -> Tensor:
        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
        alpha = ext_quantization.act_quantized_leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = self.dropout_module(alpha)
        return x_j * alpha.unsqueeze(-1)
exact/build/lib.linux-x86_64-cpython-38/exact/gatconv.py
saintslab/i-Exact
from torch._C import qscheme
from torch_geometric.nn.conv import GCNConv, SAGEConv, GCN2Conv, GraphConv, GENConv
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from torch_geometric.typing import Adj, OptTensor, Optional, Union, OptPairTensor, Size
from torch_geometric.utils import softmax
from torch_sparse import SparseTensor

import torch
import torch.nn.functional as F
from torch import Tensor
from torch.nn import Parameter, Sequential
from torch_scatter import scatter, scatter_softmax

from exact.ops import rp2input, input2rp
import exact.cpp_extension.quantization as ext_quantization
from .utils import compute_tensor_bytes, get_memory_usage
from .ops import qlinear, qbatch_norm, qelu, quantize_activation
from .conf import config
from .qsheme import QScheme
from .spmm import qmatmul, QMatmul

QGCNConv_layer_ct = 0
total_act_mem = 0
GPU = 0


def qaddmm(input, mat1, mat2, beta=1., alpha=1):
    return beta * input + alpha * qlinear.apply(mat1, mat2, None)



class QMLP(Sequential):
    def __init__(self, channels, norm,
                 bias=True, dropout=0.):
        m = []
        for i in range(1, len(channels)):
            m.append(QLinear(channels[i - 1], channels[i], bias))

            if i < len(channels) - 1:
                if norm and norm == 'batch':
                    m.append(QBatchNorm1d(channels[i], affine=True))
                elif norm:
                    raise NotImplementedError(
                        f'Normalization layer "{norm}" not supported.')
                m.append(QReLU())
                m.append(QDropout(dropout))
        super(QMLP, self).__init__(*m)


class QGENConv(GENConv):
    def __init__(self, *args, num_layers=2, norm='batch', **kwargs):
        super().__init__(*args, norm=norm, num_layers=num_layers, **kwargs)
        channels = [self.in_channels]
        for _ in range(num_layers - 1):
            channels.append(self.in_channels * 2)
        channels.append(self.out_channels)
        self.mlp = QMLP(channels, norm)
        self.reset_parameters()


    def message(self, x_j: Tensor, edge_attr: OptTensor) -> Tensor:
        msg = x_j if edge_attr is None else x_j + edge_attr
        return ext_quantization.act_quantized_relu(msg) + self.eps


class QGCN2Conv(GCN2Conv):
    def __init__(self, *args, **kwargs):
        super(QGCN2Conv, self).__init__(*args, **kwargs)


    def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""

        if self.normalize:
            if isinstance(edge_index, Tensor):
                cache = self._cached_edge_index
                if cache is None:
                    edge_index, edge_weight = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim), False,
                        self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_edge_index = (edge_index, edge_weight)
                else:
                    edge_index, edge_weight = cache[0], cache[1]

            elif isinstance(edge_index, SparseTensor):
                cache = self._cached_adj_t
                if cache is None:
                    edge_index = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim), False,
                        self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_adj_t = edge_index
                else:
                    edge_index = cache

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        x = self.propagate(edge_index, x=x, edge_weight=edge_weight, size=None)

        x.mul_(1 - self.alpha)
        x_0 = self.alpha * x_0[:x.size(0)]

        if self.weight2 is None:
            out = x.add_(x_0)
            out = qaddmm(out, out, self.weight1, beta=1. - self.beta,
                               alpha=self.beta)
        else:
            out = qaddmm(x, x, self.weight1, beta=1. - self.beta,
                               alpha=self.beta)
            out += qaddmm(x_0, x_0, self.weight2, beta=1. - self.beta,
                                alpha=self.beta)

        return out

    def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
            return edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        return qmatmul(adj_t, x, reduce=self.aggr)


class QSAGEConv(SAGEConv):
    def __init__(self, *args, **kwargs):
        super(QSAGEConv, self).__init__(*args, **kwargs)
        in_channels = self.in_channels
        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)
        # reset the lin_l and lin_l from nn.Linear to Qlinear
        bias = self.lin_l.bias is not None
        self.lin_l = QLinear(in_channels[0], self.out_channels, bias=bias, rp=True)
        if self.root_weight:
            self.lin_r = QLinear(in_channels[1], self.out_channels, bias=False, rp=True)
        self.msg_and_aggr_func = QMatmul(self.aggr)
        if config.single_precision:
            self.scheme = QScheme('sage')
        else:
            self.scheme = None
        self.reset_parameters()


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None) -> Tensor:
        """"""
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
            l_eq_r = True
        else:
            l_eq_r = False
        
        # To avoid wasting memory, 
        # if x is a tensor, we quantize it only once and pass it to the qmatmul and lin_r.
        if l_eq_r:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * x[0].shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(x[0], kept_acts)
            else:
                dim_reduced_input, randmat = x[0], None
            quantized = quantize_activation(dim_reduced_input, self.scheme)
        else:
            quantized, randmat = None, None
        # quantized = None
        # propagate_type: (x: OptPairTensor)
        out = self.propagate(edge_index, x=x, size=size, quantized=quantized, randmat=randmat)
        out = self.lin_l(out)

        x_r = x[1]
        if self.root_weight and x_r is not None:
            out += self.lin_r(x_r, quantized, randmat)

        if self.normalize:
            out = F.normalize(out, p=2., dim=-1)

        return out

    def message_and_aggregate(self, adj_t, x, quantized=None, randmat=None) -> Tensor:
        # adj_t = adj_t.set_value(None, layout=None)
        return self.msg_and_aggr_func(adj_t, x[0], quantized=quantized, randmat=randmat)


class QGraphConv(GraphConv):
    def __init__(self, *args, **kwargs):
        super(QGraphConv, self).__init__(*args, **kwargs)
        in_channels = self.in_channels
        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)
        bias = self.lin_l.bias is not None
        self.lin_l = QLinear(in_channels[0], self.out_channels, bias=bias, rp=True)
        self.lin_r = QLinear(in_channels[1], self.out_channels, bias=False, rp=True)
        self.msg_and_aggr_func = QMatmul(self.aggr)
        self.reset_parameters()


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                edge_weight: OptTensor = None, size: Size = None) -> Tensor:
        """"""
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
        # propagate_type: (x: OptPairTensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=size)
        out = self.lin_l(out)

        x_r = x[1]
        if x_r is not None:
            out += self.lin_r(x_r)
        return out


    def message_and_aggregate(self, adj_t: SparseTensor,
                              x: OptPairTensor) -> Tensor:
        return self.msg_and_aggr_func(adj_t, x[0])


class QGCNConv(GCNConv):
    def __init__(self, *args, **kwargs):
        super(QGCNConv, self).__init__(*args, **kwargs)
        self.weight = Parameter(torch.Tensor(self.out_channels, self.in_channels))
        self.msg_and_aggr_func = QMatmul(self.aggr)
        if config.single_precision:
            self.lin_scheme = QScheme('linear')
        else:
            self.lin_scheme = None
        self.reset_parameters()

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""
        if self.normalize:
            if isinstance(edge_index, Tensor):
                cache = self._cached_edge_index
                if cache is None:
                    edge_index, edge_weight = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim),
                        self.improved, self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_edge_index = (edge_index, edge_weight)
                else:
                    edge_index, edge_weight = cache[0], cache[1]

            elif isinstance(edge_index, SparseTensor):
                cache = self._cached_adj_t
                if cache is None:
                    edge_index = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim),
                        self.improved, self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_adj_t = edge_index
                else:
                    edge_index = cache
        if self.training:
            x = qlinear.apply(x, self.weight, None, None, self.lin_scheme)
        else:
            x = F.linear(x, self.weight, None)

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=None)

        if self.bias is not None:
            out += self.bias
        return out

    def message_and_aggregate(self, adj_t, x):
        return self.msg_and_aggr_func(adj_t, x)


class QLinear(torch.nn.Linear):
    def __init__(self, input_features, output_features, bias=True, rp=True):
        super(QLinear, self).__init__(input_features, output_features, bias)
        if config.single_precision:
            self.scheme = QScheme('linear')
        else:
            self.scheme = None
        self.rp = rp

    def forward(self, input, quantized=None, randmat=None):
        if self.training:
            return qlinear.apply(input, self.weight, self.bias, quantized, randmat, self.scheme, self.rp)
        else:
            return super(QLinear, self).forward(input)


class QReLU(torch.nn.Module):
    def __init__(self, inplace=False):
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_relu(input)
        else:
            return F.relu(input)


class QLeakyReLU(torch.nn.Module):
    def __init__(self, negative_slope, inplace=False):
        super().__init__()
        self.negative_slope = negative_slope

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return ext_quantization.act_quantized_leaky_relu(input, self.negative_slope)


class QELU(torch.nn.Module):
    def __init__(self, alpha, inplace=False):
        super().__init__()
        self.alpha = alpha

    def forward(self, input):
        if self.training:
            return qelu.apply(input, self.alpha)
        else:
            return F.elu(input, self.alpha)


class QDropout(torch.nn.Module):
    def __init__(self, p):
        self.p = p
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_dropout(input, self.p, self.training)
        else:
            return F.dropout(input, self.p, self.training)


class QDropout2(torch.nn.Module):
    def __init__(self, p):
        self.p = p
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_dropout(input, self.p, self.training) * (1 - self.p)
        else:
            return F.dropout(input, self.p, self.training) * (1 - self.p)


class QBatchNorm1d(torch.nn.BatchNorm1d):
    def __init__(self, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
        super(QBatchNorm1d, self).__init__(num_features, eps, momentum, affine, track_running_stats)
        if config.single_precision:
            self.scheme = QScheme('batchnorm')
        else:
            self.scheme = None

    def forward(self, input):
        if not self.training:
            return super(QBatchNorm1d, self).forward(input)

        self._check_input_dim(input)

        # exponential_average_factor is set to self.momentum
        # (when it is available) only so that it gets updated
        # in ONNX graph when this node is exported to ONNX.
        if self.momentum is None:
            exponential_average_factor = 0.0
        else:
            exponential_average_factor = self.momentum

        if self.training and self.track_running_stats:
            # TODO: if statement only here to tell the jit to skip emitting this when it is None
            if self.num_batches_tracked is not None:
                self.num_batches_tracked = self.num_batches_tracked + 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        """ Decide whether the mini-batch stats should be used for normalization rather than the buffers.
        Mini-batch stats are used in training mode, and in eval mode when buffers are None.
        """
        if self.training:
            bn_training = True
        else:
            bn_training = (self.running_mean is None) and (self.running_var is None)

        """Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
        used for normalization (i.e. in eval mode when buffers are not None).
        """
        return qbatch_norm.apply(
            input,
            # If buffers are not to be tracked, ensure that they won't be updated
            self.running_mean if not self.training or self.track_running_stats else None,
            self.running_var if not self.training or self.track_running_stats else None,
            self.weight, self.bias, bn_training, exponential_average_factor, self.eps, self.scheme)

exact/build/lib.linux-x86_64-cpython-38/exact/layers.py
saintslab/i-Exact
from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict
from collections import OrderedDict
from torch_geometric.nn.conv.gen_conv import GENConv
from torch_geometric.nn.conv.graph_conv import GraphConv
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor, device, dtype

from torch_geometric.nn.conv import GCNConv, SAGEConv, GCN2Conv, SGConv
from .layers import QBatchNorm1d, QLinear, QReLU, QGCNConv, QDropout, QDropout2, QSAGEConv, QGENConv, QGraphConv, QGCN2Conv
from .conf import config
from .gatconv import CustomGATConv, QCustomGATConv


class QModule(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        QModule.convert_layers(model)

    @staticmethod
    def convert_layers(module):
        for name, child in module.named_children():
            # Do not convert layers that are already quantized
            if isinstance(child, (QBatchNorm1d, QLinear, QReLU, QGCNConv, QGCN2Conv, QDropout, QSAGEConv, QGENConv, QCustomGATConv)):
                continue
            if isinstance(child, nn.BatchNorm1d) and config.enable_quantized_bn:
                setattr(module, name, QBatchNorm1d(child.num_features, child.eps, child.momentum,
                    child.affine, child.track_running_stats))
            elif isinstance(child, nn.Linear):
                setattr(module, name, QLinear(child.in_features, child.out_features,
                    child.bias is not None))
            elif isinstance(child, nn.ReLU):
                setattr(module, name, QReLU())
            elif isinstance(child, nn.Dropout):
                if config.dropout2:
                    setattr(module, name, QDropout2(child.p))
                else:
                    setattr(module, name, QDropout(child.p))
            elif isinstance(child, GCNConv):
                setattr(module, name, QGCNConv(child.in_channels, child.out_channels, child.improved, child.cached,
                                               child.add_self_loops, child.normalize, child.bias is not None,
                                               aggr=child.aggr))
            elif isinstance(child, GCN2Conv):
                beta = child.beta
                shared_weights = child.weight2 is None
                setattr(module, name, QGCN2Conv(child.channels, alpha=child.alpha, theta=None, layer=None, shared_weights=shared_weights,
                                                cached=child.cached, add_self_loops=child.add_self_loops, normalize=child.normalize))
                curconv = getattr(module, name)
                curconv.beta = child.beta
            elif isinstance(child, SAGEConv):
                setattr(module, name, QSAGEConv(child.in_channels, child.out_channels, child.normalize, child.root_weight,
                                               child.lin_l.bias is not None))
            elif isinstance(child, GraphConv):
                setattr(module, name, QGraphConv(child.in_channels, child.out_channels, child.aggr, child.lin_l.bias is not None))
            elif isinstance(child, CustomGATConv):
                setattr(module, name, QCustomGATConv(child.in_channels, child.out_channels, child.heads, child.concat, 
                                                     child.negative_slope, child.dropout, child.add_self_loops, child.bias is not None, 
                                                     child.residual, child.use_attn_dst))
            elif isinstance(child, GENConv):
                msg_norm = child.msg_norm is not None
                learn_msg_scale = True if (msg_norm and child.msg_norm.scale.requires_grad) else False
                learn_p = isinstance(child.p, torch.nn.Parameter)
                is_softmax = child.aggr == 'softmax'
                if is_softmax and isinstance(child.t, torch.nn.Parameter):
                    learn_t = True
                else:
                    learn_t = False
                num_layers = 0
                norm = 'batch'
                for m in child.mlp:
                    if isinstance(m, torch.nn.Linear):
                        num_layers += 1
                    if isinstance(m, (nn.BatchNorm1d, nn.LayerNorm, nn.InstanceNorm1d)):
                        if isinstance(m, nn.BatchNorm1d):
                            pass
                        elif isinstance(m, nn.LayerNorm):
                            norm = 'layer'
                        elif isinstance(m, nn.InstanceNorm1d):
                            norm = 'instance'

                setattr(module, name, QGENConv(child.in_channels, child.out_channels, child.aggr, child.initial_t, learn_t, child.initial_p, learn_p, 
                                               msg_norm, learn_msg_scale, norm, num_layers, child.eps))
            else:
                QModule.convert_layers(child)

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor]],
                        strict: bool = True):
        # remove the prefix "model." added by this wrapper
        new_state_dict = OrderedDict([("model." + k,  v) for k, v in state_dict.items()])
        return super().load_state_dict(new_state_dict, strict)

    def state_dict(self, destination=None, prefix='', keep_vars=False):
        ret = super().state_dict(destination, prefix, keep_vars)

        # remove the prefix "model." added by this wrapper
        ret = OrderedDict([(k[6:], v) for k, v in ret.items()])
        return ret

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


    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        return self.model.mini_inference(x_all, loader)
exact/build/lib.linux-x86_64-cpython-38/exact/module.py
saintslab/i-Exact
import pdb
import numpy as np
import os
import csv

import torch
import torch.nn.functional as F
from bitsandbytes import functional as bnb_F
from torch.autograd.function import Function
from torch.cuda.amp import custom_fwd, custom_bwd

from torch_sparse import matmul
from .conf import config
import exact.cpp_extension.spmm as spmm 

import exact.cpp_extension.backward_func as ext_backward_func
import exact.cpp_extension.quantization as ext_quantization

from exact.utils import empty_cache, get_memory_usage, compute_tensor_bytes, swap_to_cpu, cast_low_bit_int


def quantize_and_pack(data, bits, mn, mx):
    if config.simulate:
        N = data.shape[0]
        output = data  
        B = 2 ** bits - 1
        mn = mn - 1e-6
        mx = mx + 1e-6
        scale = B / (mx - mn) 
        output = (output - mn.view(-1, 1)) * scale.view(-1, 1)
        if config.stochastic:
            noise = output.new(output.shape).uniform_(-0.5, 0.5)
            output.add_(noise)

        output = F.relu(output)
        output = output.round_().int()
    else:
        # Pack to bitstream
        assert type(bits) == int
        pack_func = ext_quantization.pack_single_precision
        scale = (2 ** bits - 1) / (mx - mn) 
        if config.use_optimal_lo:
            dim = data.size(-1)
            a, b, c, d, f = 0.83162565, 0.02993653, 116.05528511, 0.00001225, -0.00171617
            config.lo = a + b * np.log(dim * c) + d * dim + f * dim**0.5
        lo, hi = config.lo, 3 - config.lo
        output = pack_func(data, mn, mx, scale.to(data.dtype), bits, config.stochastic, lo, hi)
        if config.swap:
            output = swap_to_cpu(output)
    return output, scale, (lo, hi)


def dequantize_and_unpack(data, shape, bits, scale, mn, lo, hi):
    if config.simulate:
        data = data / scale.view(-1, 1) + mn.view(-1, 1)
    else:
        if config.swap:
            data = data.cuda()
        N = shape[0]
        num_features = int(np.prod(shape[1:]))
        # Unpack bitstream
        assert type(bits) == int
        unpack_func = ext_quantization.unpack_single_precision
        data = unpack_func(data, bits, scale, mn, N, num_features, lo, hi)
    return data


def no_scheme_compute_quantization_bits(input):
    N = input.shape[0]
    input_flatten = input.view(N, -1)

    mn, mx = torch.min(input_flatten, 1)[0], torch.max(input_flatten, 1)[0]

    b = config.activation_compression_bits[0]
    return input_flatten, b, mn, mx


def closest_multiple(col_size, multiple):
    remainder = col_size % multiple
    if remainder > multiple / 2:
        return col_size + (multiple - remainder)
    else:
        return col_size - remainder


def quantize_activation(input, scheme):
    #input = expand_middle(input)
    # col_size is actually col_size multiplier
    col_size = config.col_size * input.size(-1)
    col_size = np.clip(col_size, input.size(-1), 1024)
    col_size = closest_multiple(col_size, input.size(-1))
    
    # Calculate the size of the reshaped tensor
    elements = input.numel()
    remaining_elements = elements % col_size
    main_part_rows = (elements - remaining_elements) // input.size(1)

    orig_feat_dim = input.size(-1)
    input2 = input[main_part_rows:].clone()

    #input = input.narrow(0, 0, main_part_rows).view(-1, col_size)
    input = input[:main_part_rows].view(-1, col_size)

    if not config.compress_activation:
        if config.swap:
            input = swap_to_cpu(input)
        return input, None, None, None

    if scheme:
        input_groups, q_bits, q_min, mx = scheme.compute_quantization_bits(input)
    else:
        input_groups, q_bits, q_min, mx = no_scheme_compute_quantization_bits(input)

    q_input, q_scale, (lo, hi) = quantize_and_pack(input_groups, q_bits, q_min, mx)
    if input.dtype == torch.float32:
        return q_input, q_bits, q_scale.to(torch.bfloat16), q_min.to(torch.bfloat16), input2, input.shape, orig_feat_dim, (lo, hi)
    else:
        return q_input, q_bits, q_scale, q_min, input2, input.shape, orig_feat_dim, (lo, hi)


def dequantize_activation(quantized, q_input_shape):
    if not config.compress_activation:
        ret = quantized[0]
        if config.swap:
            ret = ret.cuda(non_blocking=True)
        return ret

    q_input, q_bits, q_scale, q_min, input2, q_input_shape, orig_feat_dim, (lo, hi) = quantized
    if q_scale.dtype == torch.bfloat16:
        q_scale = q_scale.to(torch.float32)
        q_min = q_min.to(torch.float32)
    input = dequantize_and_unpack(q_input, q_input_shape, q_bits, q_scale, q_min, lo, hi)     

    input = input.view(-1, orig_feat_dim)
    input = torch.cat([input, input2], dim=0)
    return input.contiguous()


linear_layer_ct = 0
qmatmul_layer_ct = 0
bn_layer_ct = 0
total_act_mem = 0
GPU = 0

@torch.no_grad()
def input2rp(input, kept_acts):
    assert len(input.size()) == 2
    rand_mat_size = (input.shape[1], kept_acts)
    # Create random matrix
    def gen_rad_mat(rm_size, feat_size, device, dtype):
        bern = torch.randint(2, size=rm_size, device=device, requires_grad=False, dtype=dtype)
        return (2.0 * bern - 1) / feat_size **0.5

    rand_matrix = gen_rad_mat(rand_mat_size, kept_acts, input.device, input.dtype)
    dim_reduced_input = torch.matmul(input, rand_matrix)
    return dim_reduced_input, rand_matrix


@torch.no_grad()
def rp2input(dim_reduced_input, input_shape, rand_matrix):
    assert len(dim_reduced_input.size()) == 2
    input = torch.matmul(dim_reduced_input, rand_matrix.t())    
    return input.view(input_shape)


class qlinear(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, input, weight, bias=None, quantized=None, randmat=None, scheme=None, rp=True):
        if quantized is not None:
            if randmat is None:
                assert rp is False or config.kept_frac == 1.0
                # quantized somewhere before
                ori_input_shape, proj_input_shape = input.shape, input.shape
            else:
                assert (rp is True and config.kept_frac < 1.0) and (input.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = input.shape, torch.Size([input.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0 and rp:
                kept_acts = int(config.kept_frac * input.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(input, kept_acts)
                ori_input_shape, proj_input_shape = input.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = input, None
                ori_input_shape, proj_input_shape = input.shape, input.shape
            quantized = quantize_activation(dim_reduced_input, scheme)

        empty_cache(config.empty_cache_threshold)

        ctx.scheme = scheme
        ctx.saved = quantized, weight, bias, randmat
        ctx.other_args = ori_input_shape, proj_input_shape, rp
        res = F.linear(input, weight, bias)
        return res

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        quantized, weight, bias, randmat = ctx.saved
        ori_input_shape, q_input_shape, rp = ctx.other_args
        input = dequantize_activation(quantized, q_input_shape)
        if config.kept_frac < 1.0 and rp:
            input = rp2input(input, ori_input_shape, randmat)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_input, grad_weight, grad_bias = ext_backward_func.linear_backward(grad_output, input, weight, bias)
        # grad_input = grad_output.mm(weight)
        # grad_weight = grad_output.t().mm(input)
        # if bias is not None:
        #     grad_bias = grad_output.sum(0)
        # else:
        #     grad_bias = None
        del input, grad_output
        empty_cache(config.empty_cache_threshold)
        return grad_input, grad_weight, grad_bias, None, None, None, None



class qelu(Function):
    @staticmethod
    def forward(ctx, input, alpha, scheme=None):
        quantized = quantize_activation(input, scheme)
        empty_cache(config.empty_cache_threshold)
        ctx.scheme = scheme
        ctx.saved = quantized
        ctx.other_args = input.shape, alpha
        res = F.elu(input, alpha)
        return res

    @staticmethod
    def backward(ctx, grad_output):
        # if ctx.scheme:
        #     ctx.scheme.set_scale(grad_output)

        quantized = ctx.saved
        q_input_shape, alpha = ctx.other_args
        input = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_input = ext_backward_func._elu_backward_cuda(grad_output, input, alpha)
        return grad_input, None, None

class qbatch_norm(Function):
    @staticmethod
    def forward(ctx, input, running_mean, running_var, weight, bias,
                training, exponential_average_factor, eps, scheme):
        quantized = quantize_activation(input, scheme)

        if training:
            output, save_mean, save_var, reserve, _ = ext_backward_func._batch_norm_impl_index(input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps, True)

        else:
            output, save_mean, save_var = ext_backward_func.native_batch_norm(
                input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps)
            reserve = None

        # output, save_mean, save_var = ext_backward_func.native_batch_norm(
        #     input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps)
        # reserve = None
        ctx.scheme = scheme
        ctx.other_args = input.shape
        ctx.saved = (quantized, weight, running_mean, running_var, save_mean, save_var, training, eps, reserve)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        quantized, weight, running_mean, running_var, save_mean, save_var, training, eps, reserve = ctx.saved

        q_input_shape = ctx.other_args

        input = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved

        empty_cache(config.empty_cache_threshold)

        # if training:
        #     input = input.contiguous()
        #     grad_input, grad_weight, grad_bias = ext_backward_func.cudnn_batch_norm_backward(
        #         input, grad_output, weight, running_mean, running_var, save_mean, save_var, eps, reserve)
        # else:
        #     grad_input, grad_weight, grad_bias = ext_backward_func.native_batch_norm_backward(
        #         grad_output, input, weight, running_mean, running_var, save_mean, save_var, training, eps,
        #         [ctx.needs_input_grad[0], ctx.needs_input_grad[3], ctx.needs_input_grad[4]]
        #     )

        grad_input, grad_weight, grad_bias = ext_backward_func.native_batch_norm_backward(
            grad_output, input, weight, running_mean, running_var, save_mean, save_var, training, eps,
            [ctx.needs_input_grad[0], ctx.needs_input_grad[3], ctx.needs_input_grad[4]]
        )
        return grad_input, None, None, grad_weight, grad_bias, None, None, None, None


class qspmm_sum(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, row, rowptr, col, value, colptr, csr2csc, has_value, other, quantized=None, randmat=None, scheme=None):
        result = spmm.spmm_sum_fw(row, rowptr, col, value, colptr, csr2csc, other)
        if quantized is not None:
            if randmat is None:
                assert config.kept_frac == 1.0
                # quantized somewhere before
                ori_input_shape, proj_input_shape = other.shape, other.shape
            else:
                assert config.kept_frac < 1.0 and (other.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = other.shape, torch.Size([other.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * other.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(other, kept_acts)
                ori_input_shape, proj_input_shape = other.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = other, None
                ori_input_shape, proj_input_shape = other.shape, other.shape
            quantized = quantize_activation(dim_reduced_input, scheme)
        
        empty_cache(config.empty_cache_threshold)
        ctx.saved = row, rowptr, col, value, colptr, csr2csc, quantized, randmat
        ctx.other_args = has_value, ori_input_shape, proj_input_shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.scheme = scheme
        return result


    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        row, rowptr, col, value, colptr, csr2csc, quantized, randmat = ctx.saved
        row = col if row is None else row
        value = col if value is None else value
        colptr = col if colptr is None else colptr
        csr2csc = col if csr2csc is None else csr2csc
        has_value, ori_input_shape, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        if config.kept_frac < 1.0:
            other = rp2input(other, ori_input_shape, randmat)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_sum_bw(row, rowptr, col, value, colptr, csr2csc, other, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        del other
        empty_cache(config.empty_cache_threshold)
        return None, None, None, grad_value, None, None, None, grad_mat, None, None, None


class qspmm_mean(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        result = spmm.spmm_mean_fw(row, rowptr, col, value, rowcount, colptr, csr2csc, other)
        if quantized is not None:
            if randmat is None:
                assert config.kept_frac == 1.0
                # quantized somewhere before without random projection
                ori_input_shape, proj_input_shape = other.shape, other.shape
            else:
                assert config.kept_frac < 1.0 and (other.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = other.shape, torch.Size([other.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * other.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(other, kept_acts)
                ori_input_shape, proj_input_shape = other.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = other, None
                ori_input_shape, proj_input_shape = other.shape, other.shape
            quantized = quantize_activation(dim_reduced_input, scheme)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = row, rowptr, col, value, rowcount, colptr, csr2csc, quantized, randmat
        ctx.other_args = has_value, ori_input_shape, proj_input_shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.scheme = scheme
        return result

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        row, rowptr, col, value, rowcount, colptr, csr2csc, quantized, randmat = ctx.saved
        row = col if row is None else row
        value = col if value is None else value
        rowcount = col if rowcount is None else rowcount
        colptr = col if colptr is None else colptr
        csr2csc = col if csr2csc is None else csr2csc
        has_value, ori_input_shape, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        # here is one ugly trick: if we know value does not need gradient,
        # we actually do not need the ``other'' matrix to calculate the gradient.
        # So here we just pass a dummy matrix to the CUDA kernel.
        # TODO: engineering optimization.
        if value_requires_grad:
            other = dequantize_activation(quantized, q_input_shape)
            if config.kept_frac < 1.0:
                other = rp2input(other, ori_input_shape, randmat)
            # other = rp2input(other, ori_input_shape, randmat)

        else:
            if quantized[2] is None:
                other = torch.tensor([1.], dtype=torch.float, device='cuda')
            else:
                if quantized[2].dtype == torch.bfloat16:
                    dtype = torch.float
                else:
                    dtype = quantized[2].dtype 
                other = torch.tensor([1.], dtype=dtype, device=quantized[2].device)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_mean_bw(row, rowptr, col, value, rowcount, colptr, csr2csc, other, grad_outputs, 
                                                 has_value, value_requires_grad, mat_requires_grad)
        del other
        empty_cache(config.empty_cache_threshold)
        return None, None, None, grad_value, None, None, None, None, grad_mat, None, None, None


class qspmm_max(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, rowptr, col, value, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        output, arg_out = spmm.spmm_max_fw(rowptr, col, value, other)
        if quantized is None:
            quantized = quantize_activation(other, scheme)
        else:
            assert isinstance(quantized, tuple)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = col, value, quantized, arg_out
        ctx.other_args = has_value, other.shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.mark_non_differentiable(arg_out)
        ctx.scheme = scheme
        return output

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        col, value, quantized, arg_out = ctx.saved
        value = col if value is None else value
        has_value, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_max_bw(col, value, other, arg_out, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        return None, None, grad_value, None, grad_mat, None, None, None


class qspmm_min(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, rowptr, col, value, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        output, arg_out =  spmm.spmm_min_fw(rowptr, col, value, other)
        if quantized is None:
            quantized = quantize_activation(other, scheme)
        else:
            assert isinstance(quantized, tuple)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = col, value, quantized, arg_out
        ctx.other_args = has_value, other.shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.mark_non_differentiable(arg_out)
        ctx.scheme = scheme
        return output
   
    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        col, value, quantized, arg_out = ctx.saved
        value = col if value is None else value
        has_value, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        # if ctx.scheme:
        #     ctx.scheme.set_scale(grad_outputs)
        grad_value, grad_mat = spmm.spmm_min_bw(col, value, other, arg_out, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        return None, None, grad_value, None, grad_mat, None, None, None
exact/build/lib.linux-x86_64-cpython-38/exact/ops.py
saintslab/i-Exact
import pdb

import torch
from .conf import config
import exact.cpp_extension.calc_precision as ext_calc_precision


class QScheme(object):
    num_samples = 1
    num_matmul_layer = 0
    num_lin_layer = 0
    num_batchnorm_layer = 0
    num_sage_layer = 0
    batch = None
    update_scale = True
    layers = []
    prev_layer = None

    def __init__(self, name):
        self.initial_bits = config.initial_bits
        self.bits = config.activation_compression_bits[0]
        QScheme.layers.append(self)

        assert name in ['linear', 'matmul', 'batchnorm', 'sage']
        if name == 'linear':
            self.name = f'linear_{QScheme.num_lin_layer}'
            QScheme.num_lin_layer += 1
        elif name == 'matmul':
            self.name = f'matmul_{QScheme.num_matmul_layer}'
            QScheme.num_matmul_layer += 1
        elif name == 'batchnorm':
            self.name = f'batchnorm_{QScheme.num_batchnorm_layer}'
            QScheme.num_batchnorm_layer += 1
        elif name == 'sage':
            self.name = f'sage_{QScheme.num_sage_layer}'
            QScheme.num_sage_layer += 1
        else:
            raise ValueError


    def compute_quantization_bits(self, input):
        QScheme.prev_layer = self
        N = input.shape[0]
        input_flatten = input.view(N, -1)
        mn, mx = torch.min(input_flatten, 1)[0], torch.max(input_flatten, 1)[0]
        # range = mx - mn
        b = config.activation_compression_bits[0]
        return input_flatten, b, mn, mx
exact/build/lib.linux-x86_64-cpython-38/exact/qsheme.py
saintslab/i-Exact
import torch
from .conf import config
from .qsheme import QScheme
from .ops import qspmm_sum, qspmm_mean, qspmm_max, qspmm_min


class QMatmul(torch.nn.Module):
    def __init__(self, reduce):
        super().__init__()
        if reduce in ['sum', 'add']:
            reduce = 'sum'
        self.reduce = reduce
        if config.single_precision:
            self.scheme = QScheme('matmul')
        else:
            self.scheme = None

    def forward(self, src, other, quantized=None, randmat=None): 
        rowptr, col, value = src.csr()
        has_value = value is not None
        if self.reduce == 'sum':
            row = src.storage._row
            csr2csc = src.storage._csr2csc
            colptr = src.storage._colptr
            if has_value and value.requires_grad:
                row = src.storage.row()
            if other.requires_grad:
                row = src.storage.row()
                csr2csc = src.storage.csr2csc()
                colptr = src.storage.colptr()
            return qspmm_sum.apply(row, rowptr, col, value, colptr, csr2csc, has_value, other, 
                                quantized, randmat, self.scheme)
        elif self.reduce == 'mean':
            row = src.storage._row
            rowcount = src.storage._rowcount
            csr2csc = src.storage._csr2csc
            colptr = src.storage._colptr
            if has_value and value.requires_grad:
                row = src.storage.row()
            if other.requires_grad:
                row = src.storage.row()
                rowcount = src.storage.rowcount()
                csr2csc = src.storage.csr2csc()
                colptr = src.storage.colptr()
            return qspmm_mean.apply(row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                                    quantized, randmat, self.scheme)
        elif self.reduce == 'min':
            return qspmm_min.apply(rowptr, col, value, has_value, other, quantized, randmat, self.scheme)
        elif self.reduce == 'max':
            return qspmm_max.apply(rowptr, col, value, has_value, other, quantized, randmat, self.scheme)
        else:
            raise ValueError


def qmatmul(src, other, reduce="sum", quantized=None, randmat=None, scheme=None):
    rowptr, col, value = src.csr()
    has_value = value is not None
    if reduce in ['sum', 'add']:
        reduce = 'sum'
        row = src.storage._row
        csr2csc = src.storage._csr2csc
        colptr = src.storage._colptr
        if has_value and value.requires_grad:
            row = src.storage.row()
        if other.requires_grad:
            row = src.storage.row()
            csr2csc = src.storage.csr2csc()
            colptr = src.storage.colptr()
        return qspmm_sum.apply(row, rowptr, col, value, colptr, csr2csc, has_value, other, 
                               quantized, randmat, scheme)
    elif reduce == 'mean':
        row = src.storage._row
        rowcount = src.storage._rowcount
        csr2csc = src.storage._csr2csc
        colptr = src.storage._colptr
        if has_value and value.requires_grad:
            row = src.storage.row()
        if other.requires_grad:
            row = src.storage.row()
            rowcount = src.storage.rowcount()
            csr2csc = src.storage.csr2csc()
            colptr = src.storage.colptr()
        return qspmm_mean.apply(row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                                quantized, randmat, scheme)
    elif reduce == 'min':
        return qspmm_min.apply(rowptr, col, value, has_value, other, quantized, randmat, scheme)
    elif reduce == 'max':
        return qspmm_max.apply(rowptr, col, value, has_value, other, quantized, randmat, scheme)
    else:
        raise ValueError
exact/build/lib.linux-x86_64-cpython-38/exact/spmm.py
saintslab/i-Exact
import os
from collections import OrderedDict
import json

import torch
import numpy as np


def cast_low_bit_int(idx_vec):
    if idx_vec is None:
        return idx_vec
    assert idx_vec.dtype == torch.int64
    max_v = torch.max(idx_vec)
    if max_v <= np.iinfo(np.int16).max:
        return idx_vec.type(torch.int16)
    elif max_v < np.iinfo(np.int32).max:
        return idx_vec.type(torch.int32)
    else:
        return idx_vec


def cast_adj(data):
    data.adj_t.storage._col = cast_low_bit_int(data.adj_t.storage._col)
    data.adj_t.storage._row = cast_low_bit_int(data.adj_t.storage._row)
    data.adj_t.storage._rowptr = cast_low_bit_int(data.adj_t.storage._rowptr)
    data.adj_t.storage._rowcount = cast_low_bit_int(data.adj_t.storage._rowcount)
    data.adj_t.storage._colptr = cast_low_bit_int(data.adj_t.storage._colptr)
    data.adj_t.storage._csr2csc = cast_low_bit_int(data.adj_t.storage._csr2csc)


def swap_to_cpu(tensor):
    tensor_cpu = torch.empty(tensor.shape, dtype=tensor.dtype, device='cpu', pin_memory=True)
    tensor_cpu.copy_(tensor, non_blocking=True)
    return tensor_cpu


def get_memory_usage(gpu, print_info=False):
    """Get accurate gpu memory usage by querying torch runtime"""
    allocated = torch.cuda.memory_allocated(gpu)
    reserved = torch.cuda.memory_reserved(gpu)
    if print_info:
        print("allocated: %.2f MB" % (allocated / 1024 / 1024), flush=True)
        print("reserved:  %.2f MB" % (reserved / 1024 / 1024), flush=True)
    return allocated


def compute_tensor_bytes(tensors):
    """Compute the bytes used by a list of tensors"""
    if not isinstance(tensors, (list, tuple)):
        tensors = [tensors]

    ret = 0
    for x in tensors:
        if x.dtype in [torch.int64, torch.long]:
            ret += np.prod(x.size()) * 8
        if x.dtype in [torch.float32, torch.int, torch.int32]:
            ret += np.prod(x.size()) * 4
        elif x.dtype in [torch.bfloat16, torch.float16, torch.int16]:
            ret += np.prod(x.size()) * 2
        elif x.dtype in [torch.int8]:
            ret += np.prod(x.size())
        else:
            print(x.dtype)
            raise ValueError()
    return ret


def empty_cache(ratio):
    if ratio is None:
        return
    allocated = torch.cuda.memory_allocated(0)
    reserved = torch.cuda.memory_reserved(0)
    if reserved > 0 and allocated / reserved < ratio:
        torch.cuda.empty_cache()


def disable_cache_allocator():
    os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'


def enable_cache_allocator():
    del os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING']


class GlobalExpRecorder:
    def __init__(self):
        self.val_dict = OrderedDict()

    def record(self, key, value, float_round=6):
        if isinstance(value, (np.int32, np.int64)):
            value = int(value)
        if isinstance(value, (float, np.float32, np.float64)):
            value = round(value, float_round)

        self.val_dict[key] = value

    def dump(self, filename):
        with open(filename, "a") as fout:
            fout.write(json.dumps(self.val_dict) + '\n')
        print("Save exp results to %s" % filename)

    def clear(self):
        pass

exp_recorder = GlobalExpRecorder()

exact/build/lib.linux-x86_64-cpython-38/exact/utils.py
saintslab/i-Exact
from .utils import get_memory_usage, compute_tensor_bytes, exp_recorder, cast_adj, cast_low_bit_int
from .layers import QGCNConv, QLinear, QReLU, QBatchNorm1d, QDropout, QSAGEConv, QELU, QGCN2Conv, QGENConv
from .conf import config
from .module import QModule
from .gatconv import CustomGATConv, QCustomGATConv
exact/exact/__init__.py
saintslab/i-Exact
class QuantizationConfig:
    def __init__(self):
        self.compress_activation = True
        self.activation_compression_bits = [2]
        self.initial_bits = 8
        self.stochastic = True
        self.use_gradient = False
        self.simulate = False
        self.enable_quantized_bn = True
        self.kept_frac = 1.0
        self.single_precision = True
        self.dropout2 = False
        self.col_size = 0
        self.last_epoch = False
        self.lo = 1
        self.use_optimal_lo = False
        # Memory management flag
        self.empty_cache_threshold = None
        self.swap = False
        self.amp = False
        # Debug related flag
        self.debug_memory_op_forward = False
        self.debug_memory_op_backward = False
        self.debug_remove_bn = False
        self.debug_remove_relu = False
        self.debug_remove_dropout = False

config = QuantizationConfig()
exact/exact/conf.py
saintslab/i-Exact
from exact.layers import QDropout, QLinear
from typing import Union, Tuple, Optional
from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
                                    OptTensor)

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import Parameter, Linear
from torch_sparse import SparseTensor, set_diag
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops, softmax

from torch_geometric.nn.inits import zeros
import exact.cpp_extension.quantization as ext_quantization


class CustomGATConv(MessagePassing):
    def __init__(self, in_channels: Union[int, Tuple[int, int]],
                 out_channels: int, heads: int = 1, concat: bool = True,
                 negative_slope: float = 0.2, dropout: float = 0.0,
                 add_self_loops: bool = True, bias: bool = True, 
                 residual=True, use_attn_dst=True, **kwargs):
        kwargs.setdefault('aggr', 'add')
        super(CustomGATConv, self).__init__(node_dim=0, **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.negative_slope = negative_slope
        self.dropout = dropout
        self.add_self_loops = add_self_loops
        self.residual = residual
        self.use_attn_dst = use_attn_dst

        if isinstance(in_channels, int):
            self.lin_l = Linear(in_channels, heads * out_channels, bias=False)
            self.lin_r = self.lin_l
        else:
            self.lin_l = Linear(in_channels[0], heads * out_channels, False)
            self.lin_r = Linear(in_channels[1], heads * out_channels, False)

        if residual:
            self.res_fc = Linear(in_channels, heads * out_channels, bias=False)
        else:
            self.register_buffer("res_fc", None)

        self.att_l = Parameter(torch.Tensor(1, heads, out_channels))
        if self.use_attn_dst:
            self.att_r = Parameter(torch.Tensor(1, heads, out_channels))
        else:
            self.att_r = self.att_l
        if bias and concat:
            self.bias = Parameter(torch.Tensor(heads * out_channels))
        elif bias and not concat:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()


    def reset_parameters(self):
        gain = torch.nn.init.calculate_gain("relu")
        torch.nn.init.xavier_normal_(self.lin_l.weight, gain=gain)
        torch.nn.init.xavier_normal_(self.lin_r.weight, gain=gain)
        if isinstance(self.res_fc, Linear):
            torch.nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
        torch.nn.init.xavier_normal_(self.att_l, gain=gain)
        if self.use_attn_dst:
            torch.nn.init.xavier_normal_(self.att_r, gain=gain)
        zeros(self.bias)


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None):
        H, C = self.heads, self.out_channels

        x_l: OptTensor = None
        x_r: OptTensor = None
        alpha_l: OptTensor = None
        alpha_r: OptTensor = None
        if isinstance(x, Tensor):
            assert x.dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = x_r = self.lin_l(x).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if self.use_attn_dst:
                alpha_r = (x_r * self.att_r).sum(dim=-1)
            else:
                alpha_r = alpha_l
        else:
            x_l, x_r = x[0], x[1]
            assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.'
            x_l = self.lin_l(x_l).view(-1, H, C)
            alpha_l = (x_l * self.att_l).sum(dim=-1)
            if x_r is not None:
                x_r = self.lin_r(x_r).view(-1, H, C)
                alpha_r = (x_r * self.att_r).sum(dim=-1)

        assert x_l is not None
        assert alpha_l is not None

        if self.add_self_loops:
            if isinstance(edge_index, Tensor):
                num_nodes = x_l.size(0)
                if x_r is not None:
                    num_nodes = min(num_nodes, x_r.size(0))
                if size is not None:
                    num_nodes = min(size[0], size[1])
                edge_index, _ = remove_self_loops(edge_index)
                edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
            elif isinstance(edge_index, SparseTensor):
                edge_index = set_diag(edge_index)

        # propagate_type: (x: OptPairTensor, alpha: OptPairTensor)
        out = self.propagate(edge_index, x=(x_l, x_r),
                             alpha=(alpha_l, alpha_r), size=size)


        if self.concat:
            out = out.view(-1, self.heads * self.out_channels)
        else:
            out = out.mean(dim=1)
        if self.residual:
            out += self.res_fc(x)
        if self.bias is not None:
            out += self.bias
        return out

    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
                index: Tensor, ptr: OptTensor,
                size_i: Optional[int]) -> Tensor:
        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
        alpha = F.leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
        return x_j * alpha.unsqueeze(-1)

    def __repr__(self):
        return '{}({}, {}, heads={}, residual={}, use_attn_dst={})'.format(self.__class__.__name__,
                                                                           self.in_channels, self.out_channels, 
                                                                           self.heads, self.residual, 
                                                                           self.use_attn_dst)


class QCustomGATConv(CustomGATConv):
    def __init__(self, *args, **kwargs):
        super(QCustomGATConv, self).__init__(*args, **kwargs)
        if isinstance(self.in_channels, int):
            self.lin_l = QLinear(self.in_channels, self.heads * self.out_channels, bias=False)
            self.lin_r = self.lin_l
        else:
            self.lin_l = QLinear(self.in_channels[0], self.heads * self.out_channels, False)
            self.lin_r = QLinear(self.in_channels[1], self.heads * self.out_channels, False)

        if self.residual:
            self.res_fc = Linear(self.in_channels, self.heads * self.out_channels, bias=False)
        else:
            self.register_buffer("res_fc", None)
        self.dropout_module = QDropout(p=self.dropout)

    def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
                index: Tensor, ptr: OptTensor,
                size_i: Optional[int]) -> Tensor:
        alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
        alpha = ext_quantization.act_quantized_leaky_relu(alpha, self.negative_slope)
        alpha = softmax(alpha, index, ptr, size_i)
        self._alpha = alpha
        alpha = self.dropout_module(alpha)
        return x_j * alpha.unsqueeze(-1)
exact/exact/gatconv.py
saintslab/i-Exact
from torch._C import qscheme
from torch_geometric.nn.conv import GCNConv, SAGEConv, GCN2Conv, GraphConv, GENConv
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from torch_geometric.typing import Adj, OptTensor, Optional, Union, OptPairTensor, Size
from torch_geometric.utils import softmax
from torch_sparse import SparseTensor

import torch
import torch.nn.functional as F
from torch import Tensor
from torch.nn import Parameter, Sequential
from torch_scatter import scatter, scatter_softmax

from exact.ops import rp2input, input2rp
import exact.cpp_extension.quantization as ext_quantization
from .utils import compute_tensor_bytes, get_memory_usage
from .ops import qlinear, qbatch_norm, qelu, quantize_activation
from .conf import config
from .qsheme import QScheme
from .spmm import qmatmul, QMatmul

QGCNConv_layer_ct = 0
total_act_mem = 0
GPU = 0


def qaddmm(input, mat1, mat2, beta=1., alpha=1):
    return beta * input + alpha * qlinear.apply(mat1, mat2, None)



class QMLP(Sequential):
    def __init__(self, channels, norm,
                 bias=True, dropout=0.):
        m = []
        for i in range(1, len(channels)):
            m.append(QLinear(channels[i - 1], channels[i], bias))

            if i < len(channels) - 1:
                if norm and norm == 'batch':
                    m.append(QBatchNorm1d(channels[i], affine=True))
                elif norm:
                    raise NotImplementedError(
                        f'Normalization layer "{norm}" not supported.')
                m.append(QReLU())
                m.append(QDropout(dropout))
        super(QMLP, self).__init__(*m)


class QGENConv(GENConv):
    def __init__(self, *args, num_layers=2, norm='batch', **kwargs):
        super().__init__(*args, norm=norm, num_layers=num_layers, **kwargs)
        channels = [self.in_channels]
        for _ in range(num_layers - 1):
            channels.append(self.in_channels * 2)
        channels.append(self.out_channels)
        self.mlp = QMLP(channels, norm)
        self.reset_parameters()


    def message(self, x_j: Tensor, edge_attr: OptTensor) -> Tensor:
        msg = x_j if edge_attr is None else x_j + edge_attr
        return ext_quantization.act_quantized_relu(msg) + self.eps


class QGCN2Conv(GCN2Conv):
    def __init__(self, *args, **kwargs):
        super(QGCN2Conv, self).__init__(*args, **kwargs)


    def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""

        if self.normalize:
            if isinstance(edge_index, Tensor):
                cache = self._cached_edge_index
                if cache is None:
                    edge_index, edge_weight = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim), False,
                        self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_edge_index = (edge_index, edge_weight)
                else:
                    edge_index, edge_weight = cache[0], cache[1]

            elif isinstance(edge_index, SparseTensor):
                cache = self._cached_adj_t
                if cache is None:
                    edge_index = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim), False,
                        self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_adj_t = edge_index
                else:
                    edge_index = cache

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        x = self.propagate(edge_index, x=x, edge_weight=edge_weight, size=None)

        x.mul_(1 - self.alpha)
        x_0 = self.alpha * x_0[:x.size(0)]

        if self.weight2 is None:
            out = x.add_(x_0)
            out = qaddmm(out, out, self.weight1, beta=1. - self.beta,
                               alpha=self.beta)
        else:
            out = qaddmm(x, x, self.weight1, beta=1. - self.beta,
                               alpha=self.beta)
            out += qaddmm(x_0, x_0, self.weight2, beta=1. - self.beta,
                                alpha=self.beta)

        return out

    def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
            return edge_weight.view(-1, 1) * x_j

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        return qmatmul(adj_t, x, reduce=self.aggr)


class QSAGEConv(SAGEConv):
    def __init__(self, *args, **kwargs):
        super(QSAGEConv, self).__init__(*args, **kwargs)
        in_channels = self.in_channels
        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)
        # reset the lin_l and lin_l from nn.Linear to Qlinear
        bias = self.lin_l.bias is not None
        self.lin_l = QLinear(in_channels[0], self.out_channels, bias=bias, rp=True)
        if self.root_weight:
            self.lin_r = QLinear(in_channels[1], self.out_channels, bias=False, rp=True)
        self.msg_and_aggr_func = QMatmul(self.aggr)
        if config.single_precision:
            self.scheme = QScheme('sage')
        else:
            self.scheme = None
        self.reset_parameters()


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                size: Size = None) -> Tensor:
        """"""
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
            l_eq_r = True
        else:
            l_eq_r = False
        
        # To avoid wasting memory, 
        # if x is a tensor, we quantize it only once and pass it to the qmatmul and lin_r.
        if l_eq_r:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * x[0].shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(x[0], kept_acts)
            else:
                dim_reduced_input, randmat = x[0], None
            quantized = quantize_activation(dim_reduced_input, self.scheme)
        else:
            quantized, randmat = None, None
        # quantized = None
        # propagate_type: (x: OptPairTensor)
        out = self.propagate(edge_index, x=x, size=size, quantized=quantized, randmat=randmat)
        out = self.lin_l(out)

        x_r = x[1]
        if self.root_weight and x_r is not None:
            out += self.lin_r(x_r, quantized, randmat)

        if self.normalize:
            out = F.normalize(out, p=2., dim=-1)

        return out

    def message_and_aggregate(self, adj_t, x, quantized=None, randmat=None) -> Tensor:
        # adj_t = adj_t.set_value(None, layout=None)
        return self.msg_and_aggr_func(adj_t, x[0], quantized=quantized, randmat=randmat)


class QGraphConv(GraphConv):
    def __init__(self, *args, **kwargs):
        super(QGraphConv, self).__init__(*args, **kwargs)
        in_channels = self.in_channels
        if isinstance(in_channels, int):
            in_channels = (in_channels, in_channels)
        bias = self.lin_l.bias is not None
        self.lin_l = QLinear(in_channels[0], self.out_channels, bias=bias, rp=True)
        self.lin_r = QLinear(in_channels[1], self.out_channels, bias=False, rp=True)
        self.msg_and_aggr_func = QMatmul(self.aggr)
        self.reset_parameters()


    def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                edge_weight: OptTensor = None, size: Size = None) -> Tensor:
        """"""
        if isinstance(x, Tensor):
            x: OptPairTensor = (x, x)
        # propagate_type: (x: OptPairTensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=size)
        out = self.lin_l(out)

        x_r = x[1]
        if x_r is not None:
            out += self.lin_r(x_r)
        return out


    def message_and_aggregate(self, adj_t: SparseTensor,
                              x: OptPairTensor) -> Tensor:
        return self.msg_and_aggr_func(adj_t, x[0])


class QGCNConv(GCNConv):
    def __init__(self, *args, **kwargs):
        super(QGCNConv, self).__init__(*args, **kwargs)
        self.weight = Parameter(torch.Tensor(self.out_channels, self.in_channels))
        self.msg_and_aggr_func = QMatmul(self.aggr)
        if config.single_precision:
            self.lin_scheme = QScheme('linear')
        else:
            self.lin_scheme = None
        self.reset_parameters()

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""
        if self.normalize:
            if isinstance(edge_index, Tensor):
                cache = self._cached_edge_index
                if cache is None:
                    edge_index, edge_weight = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim),
                        self.improved, self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_edge_index = (edge_index, edge_weight)
                else:
                    edge_index, edge_weight = cache[0], cache[1]

            elif isinstance(edge_index, SparseTensor):
                cache = self._cached_adj_t
                if cache is None:
                    edge_index = gcn_norm(  # yapf: disable
                        edge_index, edge_weight, x.size(self.node_dim),
                        self.improved, self.add_self_loops, dtype=x.dtype)
                    if self.cached:
                        self._cached_adj_t = edge_index
                else:
                    edge_index = cache
        if self.training:
            x = qlinear.apply(x, self.weight, None, None, self.lin_scheme)
        else:
            x = F.linear(x, self.weight, None)

        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=None)

        if self.bias is not None:
            out += self.bias
        return out

    def message_and_aggregate(self, adj_t, x):
        return self.msg_and_aggr_func(adj_t, x)


class QLinear(torch.nn.Linear):
    def __init__(self, input_features, output_features, bias=True, rp=True):
        super(QLinear, self).__init__(input_features, output_features, bias)
        if config.single_precision:
            self.scheme = QScheme('linear')
        else:
            self.scheme = None
        self.rp = rp

    def forward(self, input, quantized=None, randmat=None):
        if self.training:
            return qlinear.apply(input, self.weight, self.bias, quantized, randmat, self.scheme, self.rp)
        else:
            return super(QLinear, self).forward(input)


class QReLU(torch.nn.Module):
    def __init__(self, inplace=False):
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_relu(input)
        else:
            return F.relu(input)


class QLeakyReLU(torch.nn.Module):
    def __init__(self, negative_slope, inplace=False):
        super().__init__()
        self.negative_slope = negative_slope

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return ext_quantization.act_quantized_leaky_relu(input, self.negative_slope)


class QELU(torch.nn.Module):
    def __init__(self, alpha, inplace=False):
        super().__init__()
        self.alpha = alpha

    def forward(self, input):
        if self.training:
            return qelu.apply(input, self.alpha)
        else:
            return F.elu(input, self.alpha)


class QDropout(torch.nn.Module):
    def __init__(self, p):
        self.p = p
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_dropout(input, self.p, self.training)
        else:
            return F.dropout(input, self.p, self.training)


class QDropout2(torch.nn.Module):
    def __init__(self, p):
        self.p = p
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.training:
            return ext_quantization.act_quantized_dropout(input, self.p, self.training) * (1 - self.p)
        else:
            return F.dropout(input, self.p, self.training) * (1 - self.p)


class QBatchNorm1d(torch.nn.BatchNorm1d):
    def __init__(self, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
        super(QBatchNorm1d, self).__init__(num_features, eps, momentum, affine, track_running_stats)
        if config.single_precision:
            self.scheme = QScheme('batchnorm')
        else:
            self.scheme = None

    def forward(self, input):
        if not self.training:
            return super(QBatchNorm1d, self).forward(input)

        self._check_input_dim(input)

        # exponential_average_factor is set to self.momentum
        # (when it is available) only so that it gets updated
        # in ONNX graph when this node is exported to ONNX.
        if self.momentum is None:
            exponential_average_factor = 0.0
        else:
            exponential_average_factor = self.momentum

        if self.training and self.track_running_stats:
            # TODO: if statement only here to tell the jit to skip emitting this when it is None
            if self.num_batches_tracked is not None:
                self.num_batches_tracked = self.num_batches_tracked + 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        """ Decide whether the mini-batch stats should be used for normalization rather than the buffers.
        Mini-batch stats are used in training mode, and in eval mode when buffers are None.
        """
        if self.training:
            bn_training = True
        else:
            bn_training = (self.running_mean is None) and (self.running_var is None)

        """Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
        passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
        used for normalization (i.e. in eval mode when buffers are not None).
        """
        return qbatch_norm.apply(
            input,
            # If buffers are not to be tracked, ensure that they won't be updated
            self.running_mean if not self.training or self.track_running_stats else None,
            self.running_var if not self.training or self.track_running_stats else None,
            self.weight, self.bias, bn_training, exponential_average_factor, self.eps, self.scheme)

exact/exact/layers.py
saintslab/i-Exact
from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict
from collections import OrderedDict
from torch_geometric.nn.conv.gen_conv import GENConv
from torch_geometric.nn.conv.graph_conv import GraphConv
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor, device, dtype

from torch_geometric.nn.conv import GCNConv, SAGEConv, GCN2Conv, SGConv
from .layers import QBatchNorm1d, QLinear, QReLU, QGCNConv, QDropout, QDropout2, QSAGEConv, QGENConv, QGraphConv, QGCN2Conv
from .conf import config
from .gatconv import CustomGATConv, QCustomGATConv


class QModule(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        QModule.convert_layers(model)

    @staticmethod
    def convert_layers(module):
        for name, child in module.named_children():
            # Do not convert layers that are already quantized
            if isinstance(child, (QBatchNorm1d, QLinear, QReLU, QGCNConv, QGCN2Conv, QDropout, QSAGEConv, QGENConv, QCustomGATConv)):
                continue
            if isinstance(child, nn.BatchNorm1d) and config.enable_quantized_bn:
                setattr(module, name, QBatchNorm1d(child.num_features, child.eps, child.momentum,
                    child.affine, child.track_running_stats))
            elif isinstance(child, nn.Linear):
                setattr(module, name, QLinear(child.in_features, child.out_features,
                    child.bias is not None))
            elif isinstance(child, nn.ReLU):
                setattr(module, name, QReLU())
            elif isinstance(child, nn.Dropout):
                if config.dropout2:
                    setattr(module, name, QDropout2(child.p))
                else:
                    setattr(module, name, QDropout(child.p))
            elif isinstance(child, GCNConv):
                setattr(module, name, QGCNConv(child.in_channels, child.out_channels, child.improved, child.cached,
                                               child.add_self_loops, child.normalize, child.bias is not None,
                                               aggr=child.aggr))
            elif isinstance(child, GCN2Conv):
                beta = child.beta
                shared_weights = child.weight2 is None
                setattr(module, name, QGCN2Conv(child.channels, alpha=child.alpha, theta=None, layer=None, shared_weights=shared_weights,
                                                cached=child.cached, add_self_loops=child.add_self_loops, normalize=child.normalize))
                curconv = getattr(module, name)
                curconv.beta = child.beta
            elif isinstance(child, SAGEConv):
                setattr(module, name, QSAGEConv(child.in_channels, child.out_channels, child.normalize, child.root_weight,
                                               child.lin_l.bias is not None))
            elif isinstance(child, GraphConv):
                setattr(module, name, QGraphConv(child.in_channels, child.out_channels, child.aggr, child.lin_l.bias is not None))
            elif isinstance(child, CustomGATConv):
                setattr(module, name, QCustomGATConv(child.in_channels, child.out_channels, child.heads, child.concat, 
                                                     child.negative_slope, child.dropout, child.add_self_loops, child.bias is not None, 
                                                     child.residual, child.use_attn_dst))
            elif isinstance(child, GENConv):
                msg_norm = child.msg_norm is not None
                learn_msg_scale = True if (msg_norm and child.msg_norm.scale.requires_grad) else False
                learn_p = isinstance(child.p, torch.nn.Parameter)
                is_softmax = child.aggr == 'softmax'
                if is_softmax and isinstance(child.t, torch.nn.Parameter):
                    learn_t = True
                else:
                    learn_t = False
                num_layers = 0
                norm = 'batch'
                for m in child.mlp:
                    if isinstance(m, torch.nn.Linear):
                        num_layers += 1
                    if isinstance(m, (nn.BatchNorm1d, nn.LayerNorm, nn.InstanceNorm1d)):
                        if isinstance(m, nn.BatchNorm1d):
                            pass
                        elif isinstance(m, nn.LayerNorm):
                            norm = 'layer'
                        elif isinstance(m, nn.InstanceNorm1d):
                            norm = 'instance'

                setattr(module, name, QGENConv(child.in_channels, child.out_channels, child.aggr, child.initial_t, learn_t, child.initial_p, learn_p, 
                                               msg_norm, learn_msg_scale, norm, num_layers, child.eps))
            else:
                QModule.convert_layers(child)

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def load_state_dict(self, state_dict: Union[Dict[str, Tensor], Dict[str, Tensor]],
                        strict: bool = True):
        # remove the prefix "model." added by this wrapper
        new_state_dict = OrderedDict([("model." + k,  v) for k, v in state_dict.items()])
        return super().load_state_dict(new_state_dict, strict)

    def state_dict(self, destination=None, prefix='', keep_vars=False):
        ret = super().state_dict(destination, prefix, keep_vars)

        # remove the prefix "model." added by this wrapper
        ret = OrderedDict([(k[6:], v) for k, v in ret.items()])
        return ret

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


    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        return self.model.mini_inference(x_all, loader)
exact/exact/module.py
saintslab/i-Exact
import pdb
import numpy as np
import os
import csv

import torch
import torch.nn.functional as F
from bitsandbytes import functional as bnb_F
from torch.autograd.function import Function
from torch.cuda.amp import custom_fwd, custom_bwd

from torch_sparse import matmul
from .conf import config
import exact.cpp_extension.spmm as spmm 

import exact.cpp_extension.backward_func as ext_backward_func
import exact.cpp_extension.quantization as ext_quantization

from exact.utils import empty_cache, get_memory_usage, compute_tensor_bytes, swap_to_cpu, cast_low_bit_int


def quantize_and_pack(data, bits, mn, mx):
    if config.simulate:
        N = data.shape[0]
        output = data  
        B = 2 ** bits - 1
        mn = mn - 1e-6
        mx = mx + 1e-6
        scale = B / (mx - mn) 
        output = (output - mn.view(-1, 1)) * scale.view(-1, 1)
        if config.stochastic:
            noise = output.new(output.shape).uniform_(-0.5, 0.5)
            output.add_(noise)

        output = F.relu(output)
        output = output.round_().int()
    else:
        # Pack to bitstream
        assert type(bits) == int
        pack_func = ext_quantization.pack_single_precision
        scale = (2 ** bits - 1) / (mx - mn) 
        if config.use_optimal_lo:
            dim = data.size(-1)
            a, b, c, d, f = 0.83162565, 0.02993653, 116.05528511, 0.00001225, -0.00171617
            config.lo = a + b * np.log(dim * c) + d * dim + f * dim**0.5
        lo, hi = config.lo, 3 - config.lo
        output = pack_func(data, mn, mx, scale.to(data.dtype), bits, config.stochastic, lo, hi)
        if config.swap:
            output = swap_to_cpu(output)
    return output, scale, (lo, hi)


def dequantize_and_unpack(data, shape, bits, scale, mn, lo, hi):
    if config.simulate:
        data = data / scale.view(-1, 1) + mn.view(-1, 1)
    else:
        if config.swap:
            data = data.cuda()
        N = shape[0]
        num_features = int(np.prod(shape[1:]))
        # Unpack bitstream
        assert type(bits) == int
        unpack_func = ext_quantization.unpack_single_precision
        data = unpack_func(data, bits, scale, mn, N, num_features, lo, hi)
    return data


def no_scheme_compute_quantization_bits(input):
    N = input.shape[0]
    input_flatten = input.view(N, -1)

    mn, mx = torch.min(input_flatten, 1)[0], torch.max(input_flatten, 1)[0]

    b = config.activation_compression_bits[0]
    return input_flatten, b, mn, mx


def closest_multiple(col_size, multiple):
    remainder = col_size % multiple
    if remainder > multiple / 2:
        return col_size + (multiple - remainder)
    else:
        return col_size - remainder


def quantize_activation(input, scheme):
    #input = expand_middle(input)
    # col_size is actually col_size multiplier
    col_size = config.col_size * input.size(-1)
    col_size = np.clip(col_size, input.size(-1), 1024)
    col_size = closest_multiple(col_size, input.size(-1))
    
    # Calculate the size of the reshaped tensor
    elements = input.numel()
    remaining_elements = elements % col_size
    main_part_rows = (elements - remaining_elements) // input.size(1)

    orig_feat_dim = input.size(-1)
    input2 = input[main_part_rows:].clone()

    #input = input.narrow(0, 0, main_part_rows).view(-1, col_size)
    input = input[:main_part_rows].view(-1, col_size)

    if not config.compress_activation:
        if config.swap:
            input = swap_to_cpu(input)
        return input, None, None, None

    if scheme:
        input_groups, q_bits, q_min, mx = scheme.compute_quantization_bits(input)
    else:
        input_groups, q_bits, q_min, mx = no_scheme_compute_quantization_bits(input)

    q_input, q_scale, (lo, hi) = quantize_and_pack(input_groups, q_bits, q_min, mx)
    if input.dtype == torch.float32:
        return q_input, q_bits, q_scale.to(torch.bfloat16), q_min.to(torch.bfloat16), input2, input.shape, orig_feat_dim, (lo, hi)
    else:
        return q_input, q_bits, q_scale, q_min, input2, input.shape, orig_feat_dim, (lo, hi)


def dequantize_activation(quantized, q_input_shape):
    if not config.compress_activation:
        ret = quantized[0]
        if config.swap:
            ret = ret.cuda(non_blocking=True)
        return ret

    q_input, q_bits, q_scale, q_min, input2, q_input_shape, orig_feat_dim, (lo, hi) = quantized
    if q_scale.dtype == torch.bfloat16:
        q_scale = q_scale.to(torch.float32)
        q_min = q_min.to(torch.float32)
    input = dequantize_and_unpack(q_input, q_input_shape, q_bits, q_scale, q_min, lo, hi)     

    input = input.view(-1, orig_feat_dim)
    input = torch.cat([input, input2], dim=0)
    return input.contiguous()


linear_layer_ct = 0
qmatmul_layer_ct = 0
bn_layer_ct = 0
total_act_mem = 0
GPU = 0

@torch.no_grad()
def input2rp(input, kept_acts):
    assert len(input.size()) == 2
    rand_mat_size = (input.shape[1], kept_acts)
    # Create random matrix
    def gen_rad_mat(rm_size, feat_size, device, dtype):
        bern = torch.randint(2, size=rm_size, device=device, requires_grad=False, dtype=dtype)
        return (2.0 * bern - 1) / feat_size **0.5

    rand_matrix = gen_rad_mat(rand_mat_size, kept_acts, input.device, input.dtype)
    dim_reduced_input = torch.matmul(input, rand_matrix)
    return dim_reduced_input, rand_matrix


@torch.no_grad()
def rp2input(dim_reduced_input, input_shape, rand_matrix):
    assert len(dim_reduced_input.size()) == 2
    input = torch.matmul(dim_reduced_input, rand_matrix.t())    
    return input.view(input_shape)


class qlinear(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, input, weight, bias=None, quantized=None, randmat=None, scheme=None, rp=True):
        if quantized is not None:
            if randmat is None:
                assert rp is False or config.kept_frac == 1.0
                # quantized somewhere before
                ori_input_shape, proj_input_shape = input.shape, input.shape
            else:
                assert (rp is True and config.kept_frac < 1.0) and (input.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = input.shape, torch.Size([input.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0 and rp:
                kept_acts = int(config.kept_frac * input.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(input, kept_acts)
                ori_input_shape, proj_input_shape = input.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = input, None
                ori_input_shape, proj_input_shape = input.shape, input.shape
            quantized = quantize_activation(dim_reduced_input, scheme)

        empty_cache(config.empty_cache_threshold)

        ctx.scheme = scheme
        ctx.saved = quantized, weight, bias, randmat
        ctx.other_args = ori_input_shape, proj_input_shape, rp
        res = F.linear(input, weight, bias)
        return res

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        quantized, weight, bias, randmat = ctx.saved
        ori_input_shape, q_input_shape, rp = ctx.other_args
        input = dequantize_activation(quantized, q_input_shape)
        if config.kept_frac < 1.0 and rp:
            input = rp2input(input, ori_input_shape, randmat)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_input, grad_weight, grad_bias = ext_backward_func.linear_backward(grad_output, input, weight, bias)
        # grad_input = grad_output.mm(weight)
        # grad_weight = grad_output.t().mm(input)
        # if bias is not None:
        #     grad_bias = grad_output.sum(0)
        # else:
        #     grad_bias = None
        del input, grad_output
        empty_cache(config.empty_cache_threshold)
        return grad_input, grad_weight, grad_bias, None, None, None, None



class qelu(Function):
    @staticmethod
    def forward(ctx, input, alpha, scheme=None):
        quantized = quantize_activation(input, scheme)
        empty_cache(config.empty_cache_threshold)
        ctx.scheme = scheme
        ctx.saved = quantized
        ctx.other_args = input.shape, alpha
        res = F.elu(input, alpha)
        return res

    @staticmethod
    def backward(ctx, grad_output):
        # if ctx.scheme:
        #     ctx.scheme.set_scale(grad_output)

        quantized = ctx.saved
        q_input_shape, alpha = ctx.other_args
        input = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_input = ext_backward_func._elu_backward_cuda(grad_output, input, alpha)
        return grad_input, None, None

class qbatch_norm(Function):
    @staticmethod
    def forward(ctx, input, running_mean, running_var, weight, bias,
                training, exponential_average_factor, eps, scheme):
        quantized = quantize_activation(input, scheme)

        if training:
            output, save_mean, save_var, reserve, _ = ext_backward_func._batch_norm_impl_index(input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps, True)

        else:
            output, save_mean, save_var = ext_backward_func.native_batch_norm(
                input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps)
            reserve = None

        # output, save_mean, save_var = ext_backward_func.native_batch_norm(
        #     input, weight, bias, running_mean, running_var, training, exponential_average_factor, eps)
        # reserve = None
        ctx.scheme = scheme
        ctx.other_args = input.shape
        ctx.saved = (quantized, weight, running_mean, running_var, save_mean, save_var, training, eps, reserve)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        quantized, weight, running_mean, running_var, save_mean, save_var, training, eps, reserve = ctx.saved

        q_input_shape = ctx.other_args

        input = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved

        empty_cache(config.empty_cache_threshold)

        # if training:
        #     input = input.contiguous()
        #     grad_input, grad_weight, grad_bias = ext_backward_func.cudnn_batch_norm_backward(
        #         input, grad_output, weight, running_mean, running_var, save_mean, save_var, eps, reserve)
        # else:
        #     grad_input, grad_weight, grad_bias = ext_backward_func.native_batch_norm_backward(
        #         grad_output, input, weight, running_mean, running_var, save_mean, save_var, training, eps,
        #         [ctx.needs_input_grad[0], ctx.needs_input_grad[3], ctx.needs_input_grad[4]]
        #     )

        grad_input, grad_weight, grad_bias = ext_backward_func.native_batch_norm_backward(
            grad_output, input, weight, running_mean, running_var, save_mean, save_var, training, eps,
            [ctx.needs_input_grad[0], ctx.needs_input_grad[3], ctx.needs_input_grad[4]]
        )
        return grad_input, None, None, grad_weight, grad_bias, None, None, None, None


class qspmm_sum(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, row, rowptr, col, value, colptr, csr2csc, has_value, other, quantized=None, randmat=None, scheme=None):
        result = spmm.spmm_sum_fw(row, rowptr, col, value, colptr, csr2csc, other)
        if quantized is not None:
            if randmat is None:
                assert config.kept_frac == 1.0
                # quantized somewhere before
                ori_input_shape, proj_input_shape = other.shape, other.shape
            else:
                assert config.kept_frac < 1.0 and (other.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = other.shape, torch.Size([other.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * other.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(other, kept_acts)
                ori_input_shape, proj_input_shape = other.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = other, None
                ori_input_shape, proj_input_shape = other.shape, other.shape
            quantized = quantize_activation(dim_reduced_input, scheme)
        
        empty_cache(config.empty_cache_threshold)
        ctx.saved = row, rowptr, col, value, colptr, csr2csc, quantized, randmat
        ctx.other_args = has_value, ori_input_shape, proj_input_shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.scheme = scheme
        return result


    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        row, rowptr, col, value, colptr, csr2csc, quantized, randmat = ctx.saved
        row = col if row is None else row
        value = col if value is None else value
        colptr = col if colptr is None else colptr
        csr2csc = col if csr2csc is None else csr2csc
        has_value, ori_input_shape, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        if config.kept_frac < 1.0:
            other = rp2input(other, ori_input_shape, randmat)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_sum_bw(row, rowptr, col, value, colptr, csr2csc, other, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        del other
        empty_cache(config.empty_cache_threshold)
        return None, None, None, grad_value, None, None, None, grad_mat, None, None, None


class qspmm_mean(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        result = spmm.spmm_mean_fw(row, rowptr, col, value, rowcount, colptr, csr2csc, other)
        if quantized is not None:
            if randmat is None:
                assert config.kept_frac == 1.0
                # quantized somewhere before without random projection
                ori_input_shape, proj_input_shape = other.shape, other.shape
            else:
                assert config.kept_frac < 1.0 and (other.shape[1] == randmat.shape[0])
                ori_input_shape, proj_input_shape = other.shape, torch.Size([other.shape[0], randmat.shape[1]])
                # this is quantized random projected data
        else:
            if config.kept_frac < 1.0:
                kept_acts = int(config.kept_frac * other.shape[1] + 0.999)
                dim_reduced_input, randmat = input2rp(other, kept_acts)
                ori_input_shape, proj_input_shape = other.shape, dim_reduced_input.shape
            else:
                dim_reduced_input, randmat = other, None
                ori_input_shape, proj_input_shape = other.shape, other.shape
            quantized = quantize_activation(dim_reduced_input, scheme)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = row, rowptr, col, value, rowcount, colptr, csr2csc, quantized, randmat
        ctx.other_args = has_value, ori_input_shape, proj_input_shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.scheme = scheme
        return result

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        row, rowptr, col, value, rowcount, colptr, csr2csc, quantized, randmat = ctx.saved
        row = col if row is None else row
        value = col if value is None else value
        rowcount = col if rowcount is None else rowcount
        colptr = col if colptr is None else colptr
        csr2csc = col if csr2csc is None else csr2csc
        has_value, ori_input_shape, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        # here is one ugly trick: if we know value does not need gradient,
        # we actually do not need the ``other'' matrix to calculate the gradient.
        # So here we just pass a dummy matrix to the CUDA kernel.
        # TODO: engineering optimization.
        if value_requires_grad:
            other = dequantize_activation(quantized, q_input_shape)
            if config.kept_frac < 1.0:
                other = rp2input(other, ori_input_shape, randmat)
            # other = rp2input(other, ori_input_shape, randmat)

        else:
            if quantized[2] is None:
                other = torch.tensor([1.], dtype=torch.float, device='cuda')
            else:
                if quantized[2].dtype == torch.bfloat16:
                    dtype = torch.float
                else:
                    dtype = quantized[2].dtype 
                other = torch.tensor([1.], dtype=dtype, device=quantized[2].device)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_mean_bw(row, rowptr, col, value, rowcount, colptr, csr2csc, other, grad_outputs, 
                                                 has_value, value_requires_grad, mat_requires_grad)
        del other
        empty_cache(config.empty_cache_threshold)
        return None, None, None, grad_value, None, None, None, None, grad_mat, None, None, None


class qspmm_max(Function):

    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, rowptr, col, value, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        output, arg_out = spmm.spmm_max_fw(rowptr, col, value, other)
        if quantized is None:
            quantized = quantize_activation(other, scheme)
        else:
            assert isinstance(quantized, tuple)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = col, value, quantized, arg_out
        ctx.other_args = has_value, other.shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.mark_non_differentiable(arg_out)
        ctx.scheme = scheme
        return output

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        col, value, quantized, arg_out = ctx.saved
        value = col if value is None else value
        has_value, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        grad_value, grad_mat = spmm.spmm_max_bw(col, value, other, arg_out, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        return None, None, grad_value, None, grad_mat, None, None, None


class qspmm_min(Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16)
    def forward(ctx, rowptr, col, value, has_value, other, 
                quantized=None, randmat=None, scheme=None):
        output, arg_out =  spmm.spmm_min_fw(rowptr, col, value, other)
        if quantized is None:
            quantized = quantize_activation(other, scheme)
        else:
            assert isinstance(quantized, tuple)
        empty_cache(config.empty_cache_threshold)
        ctx.saved = col, value, quantized, arg_out
        ctx.other_args = has_value, other.shape, value.requires_grad if has_value else False, other.requires_grad
        ctx.mark_non_differentiable(arg_out)
        ctx.scheme = scheme
        return output
   
    @staticmethod
    @custom_bwd
    def backward(ctx, grad_outputs):
        col, value, quantized, arg_out = ctx.saved
        value = col if value is None else value
        has_value, q_input_shape, value_requires_grad, mat_requires_grad = ctx.other_args
        other = dequantize_activation(quantized, q_input_shape)
        del quantized, ctx.saved
        empty_cache(config.empty_cache_threshold)
        # if ctx.scheme:
        #     ctx.scheme.set_scale(grad_outputs)
        grad_value, grad_mat = spmm.spmm_min_bw(col, value, other, arg_out, grad_outputs, 
                                                has_value, value_requires_grad, mat_requires_grad)
        return None, None, grad_value, None, grad_mat, None, None, None
exact/exact/ops.py
saintslab/i-Exact
import pdb

import torch
from .conf import config
import exact.cpp_extension.calc_precision as ext_calc_precision


class QScheme(object):
    num_samples = 1
    num_matmul_layer = 0
    num_lin_layer = 0
    num_batchnorm_layer = 0
    num_sage_layer = 0
    batch = None
    update_scale = True
    layers = []
    prev_layer = None

    def __init__(self, name):
        self.initial_bits = config.initial_bits
        self.bits = config.activation_compression_bits[0]
        QScheme.layers.append(self)

        assert name in ['linear', 'matmul', 'batchnorm', 'sage']
        if name == 'linear':
            self.name = f'linear_{QScheme.num_lin_layer}'
            QScheme.num_lin_layer += 1
        elif name == 'matmul':
            self.name = f'matmul_{QScheme.num_matmul_layer}'
            QScheme.num_matmul_layer += 1
        elif name == 'batchnorm':
            self.name = f'batchnorm_{QScheme.num_batchnorm_layer}'
            QScheme.num_batchnorm_layer += 1
        elif name == 'sage':
            self.name = f'sage_{QScheme.num_sage_layer}'
            QScheme.num_sage_layer += 1
        else:
            raise ValueError


    def compute_quantization_bits(self, input):
        QScheme.prev_layer = self
        N = input.shape[0]
        input_flatten = input.view(N, -1)
        mn, mx = torch.min(input_flatten, 1)[0], torch.max(input_flatten, 1)[0]
        # range = mx - mn
        b = config.activation_compression_bits[0]
        return input_flatten, b, mn, mx
exact/exact/qsheme.py
saintslab/i-Exact
import torch
from .conf import config
from .qsheme import QScheme
from .ops import qspmm_sum, qspmm_mean, qspmm_max, qspmm_min


class QMatmul(torch.nn.Module):
    def __init__(self, reduce):
        super().__init__()
        if reduce in ['sum', 'add']:
            reduce = 'sum'
        self.reduce = reduce
        if config.single_precision:
            self.scheme = QScheme('matmul')
        else:
            self.scheme = None

    def forward(self, src, other, quantized=None, randmat=None): 
        rowptr, col, value = src.csr()
        has_value = value is not None
        if self.reduce == 'sum':
            row = src.storage._row
            csr2csc = src.storage._csr2csc
            colptr = src.storage._colptr
            if has_value and value.requires_grad:
                row = src.storage.row()
            if other.requires_grad:
                row = src.storage.row()
                csr2csc = src.storage.csr2csc()
                colptr = src.storage.colptr()
            return qspmm_sum.apply(row, rowptr, col, value, colptr, csr2csc, has_value, other, 
                                quantized, randmat, self.scheme)
        elif self.reduce == 'mean':
            row = src.storage._row
            rowcount = src.storage._rowcount
            csr2csc = src.storage._csr2csc
            colptr = src.storage._colptr
            if has_value and value.requires_grad:
                row = src.storage.row()
            if other.requires_grad:
                row = src.storage.row()
                rowcount = src.storage.rowcount()
                csr2csc = src.storage.csr2csc()
                colptr = src.storage.colptr()
            return qspmm_mean.apply(row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                                    quantized, randmat, self.scheme)
        elif self.reduce == 'min':
            return qspmm_min.apply(rowptr, col, value, has_value, other, quantized, randmat, self.scheme)
        elif self.reduce == 'max':
            return qspmm_max.apply(rowptr, col, value, has_value, other, quantized, randmat, self.scheme)
        else:
            raise ValueError


def qmatmul(src, other, reduce="sum", quantized=None, randmat=None, scheme=None):
    rowptr, col, value = src.csr()
    has_value = value is not None
    if reduce in ['sum', 'add']:
        reduce = 'sum'
        row = src.storage._row
        csr2csc = src.storage._csr2csc
        colptr = src.storage._colptr
        if has_value and value.requires_grad:
            row = src.storage.row()
        if other.requires_grad:
            row = src.storage.row()
            csr2csc = src.storage.csr2csc()
            colptr = src.storage.colptr()
        return qspmm_sum.apply(row, rowptr, col, value, colptr, csr2csc, has_value, other, 
                               quantized, randmat, scheme)
    elif reduce == 'mean':
        row = src.storage._row
        rowcount = src.storage._rowcount
        csr2csc = src.storage._csr2csc
        colptr = src.storage._colptr
        if has_value and value.requires_grad:
            row = src.storage.row()
        if other.requires_grad:
            row = src.storage.row()
            rowcount = src.storage.rowcount()
            csr2csc = src.storage.csr2csc()
            colptr = src.storage.colptr()
        return qspmm_mean.apply(row, rowptr, col, value, rowcount, colptr, csr2csc, has_value, other, 
                                quantized, randmat, scheme)
    elif reduce == 'min':
        return qspmm_min.apply(rowptr, col, value, has_value, other, quantized, randmat, scheme)
    elif reduce == 'max':
        return qspmm_max.apply(rowptr, col, value, has_value, other, quantized, randmat, scheme)
    else:
        raise ValueError
exact/exact/spmm.py
saintslab/i-Exact
import os
from collections import OrderedDict
import json

import torch
import numpy as np


def cast_low_bit_int(idx_vec):
    if idx_vec is None:
        return idx_vec
    assert idx_vec.dtype == torch.int64
    max_v = torch.max(idx_vec)
    if max_v <= np.iinfo(np.int16).max:
        return idx_vec.type(torch.int16)
    elif max_v < np.iinfo(np.int32).max:
        return idx_vec.type(torch.int32)
    else:
        return idx_vec


def cast_adj(data):
    data.adj_t.storage._col = cast_low_bit_int(data.adj_t.storage._col)
    data.adj_t.storage._row = cast_low_bit_int(data.adj_t.storage._row)
    data.adj_t.storage._rowptr = cast_low_bit_int(data.adj_t.storage._rowptr)
    data.adj_t.storage._rowcount = cast_low_bit_int(data.adj_t.storage._rowcount)
    data.adj_t.storage._colptr = cast_low_bit_int(data.adj_t.storage._colptr)
    data.adj_t.storage._csr2csc = cast_low_bit_int(data.adj_t.storage._csr2csc)


def swap_to_cpu(tensor):
    tensor_cpu = torch.empty(tensor.shape, dtype=tensor.dtype, device='cpu', pin_memory=True)
    tensor_cpu.copy_(tensor, non_blocking=True)
    return tensor_cpu


def get_memory_usage(gpu, print_info=False):
    """Get accurate gpu memory usage by querying torch runtime"""
    allocated = torch.cuda.memory_allocated(gpu)
    reserved = torch.cuda.memory_reserved(gpu)
    if print_info:
        print("allocated: %.2f MB" % (allocated / 1024 / 1024), flush=True)
        print("reserved:  %.2f MB" % (reserved / 1024 / 1024), flush=True)
    return allocated


def compute_tensor_bytes(tensors):
    """Compute the bytes used by a list of tensors"""
    if not isinstance(tensors, (list, tuple)):
        tensors = [tensors]

    ret = 0
    for x in tensors:
        if x.dtype in [torch.int64, torch.long]:
            ret += np.prod(x.size()) * 8
        if x.dtype in [torch.float32, torch.int, torch.int32]:
            ret += np.prod(x.size()) * 4
        elif x.dtype in [torch.bfloat16, torch.float16, torch.int16]:
            ret += np.prod(x.size()) * 2
        elif x.dtype in [torch.int8]:
            ret += np.prod(x.size())
        else:
            print(x.dtype)
            raise ValueError()
    return ret


def empty_cache(ratio):
    if ratio is None:
        return
    allocated = torch.cuda.memory_allocated(0)
    reserved = torch.cuda.memory_reserved(0)
    if reserved > 0 and allocated / reserved < ratio:
        torch.cuda.empty_cache()


def disable_cache_allocator():
    os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'


def enable_cache_allocator():
    del os.environ['PYTORCH_NO_CUDA_MEMORY_CACHING']


class GlobalExpRecorder:
    def __init__(self):
        self.val_dict = OrderedDict()

    def record(self, key, value, float_round=6):
        if isinstance(value, (np.int32, np.int64)):
            value = int(value)
        if isinstance(value, (float, np.float32, np.float64)):
            value = round(value, float_round)

        self.val_dict[key] = value

    def dump(self, filename):
        with open(filename, "a") as fout:
            fout.write(json.dumps(self.val_dict) + '\n')
        print("Save exp results to %s" % filename)

    def clear(self):
        pass

exp_recorder = GlobalExpRecorder()

exact/exact/utils.py
saintslab/i-Exact
from setuptools import setup, find_packages
from torch.utils import cpp_extension

setup(name='exact',
      ext_modules=[
          cpp_extension.CUDAExtension(
              'exact.cpp_extension.calc_precision',
              ['exact/cpp_extension/calc_precision.cc']
          ),
          cpp_extension.CUDAExtension(
              'exact.cpp_extension.backward_func',
              ['exact/cpp_extension/backward_func.cc', 'exact/cpp_extension/backward_func_cuda_kernel.cu']
          ),
          cpp_extension.CUDAExtension(
              'exact.cpp_extension.quantization',
              ['exact/cpp_extension/quantization.cc', 'exact/cpp_extension/quantization_cuda_kernel.cu'],
              extra_compile_args={'nvcc': ['--expt-extended-lambda']}
          ),
          cpp_extension.CUDAExtension(
              'exact.cpp_extension.spmm',
              ['exact/cpp_extension/spmm.cc', 'exact/cpp_extension/spmm_cuda.cu']
          ),
      ],
      cmdclass={'build_ext': cpp_extension.BuildExtension},
      packages=find_packages()
)

exact/setup.py
saintslab/i-Exact
import pdb
import torch
import torch.nn.functional as F
from torch_geometric.nn.conv import GCNConv, SAGEConv
from torch_sparse import SparseTensor

from exact import config, CustomGATConv



class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers, num_heads,
                 dropout, input_drop, attn_drop, edge_drop, residual, use_attn_dst):
        super(GAT, self).__init__()
        self.convs = torch.nn.ModuleList()
        self.convs.append(CustomGATConv(in_channels, hidden_channels, num_heads, 
                                        dropout=attn_drop, add_self_loops=False, residual=residual, use_attn_dst=use_attn_dst))
        self.bns = torch.nn.ModuleList()
        self.bns.append(torch.nn.BatchNorm1d(hidden_channels * num_heads))
        for _ in range(num_layers -2):
            self.convs.append(
                CustomGATConv(hidden_channels * num_heads, hidden_channels, num_heads, dropout=attn_drop, add_self_loops=False, 
                              residual=residual, use_attn_dst=use_attn_dst))
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels * num_heads))
        self.convs.append(CustomGATConv(hidden_channels * num_heads, out_channels, 1, dropout=attn_drop, add_self_loops=False, 
                                        residual=residual, use_attn_dst=use_attn_dst))
        self.input_dropout = torch.nn.Dropout(p=input_drop)
        self.dropout = torch.nn.Dropout(p=dropout)
        self.act = torch.nn.ReLU()
        self.edge_drop = edge_drop


    def forward(self, x, adj_t):
        x = self.input_dropout(x)
        for i, conv in enumerate(self.convs[:-1]):
            if self.training and self.edge_drop > 0:
                nnz = adj_t.nnz()
                bound = int(nnz * self.edge_drop)
                perm = torch.randperm(nnz).cuda()
                eids = perm[bound:]
                row, col, value = adj_t.storage._row, adj_t.storage._col, adj_t.storage._value
                row, col, value = row[eids], col[eids], value[eids]
                adj_t = SparseTensor(row=row, col=col, value=value, 
                                     sparse_sizes=(adj_t.size(0), adj_t.size(1)), is_sorted=False)
            x = conv(x, adj_t)
            x = self.bns[i](x)
            x = self.act(x)
            if not config.debug_remove_dropout:
                x = self.dropout(x)
        x = self.convs[-1](x, adj_t)
        return x.log_softmax(dim=-1)


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if not config.debug_remove_bn:
            for bn in self.bns:
                bn.reset_parameters()


class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(GCN, self).__init__()
        self.convs = torch.nn.ModuleList()
        self.convs.append(GCNConv(in_channels, hidden_channels, cached=True, normalize=False))
        self.bns = torch.nn.ModuleList()
        self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(
                GCNConv(hidden_channels, hidden_channels, cached=True, normalize=False))
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.convs.append(GCNConv(hidden_channels, out_channels, cached=True, normalize=False))
        self.dropout = torch.nn.Dropout(p=dropout)
        self.act = torch.nn.ReLU()

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if not config.debug_remove_bn:
            for bn in self.bns:
                bn.reset_parameters()

    def forward(self, x, adj_t):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, adj_t)
            x = self.bns[i](x)
            x = self.act(x)
            if not config.debug_remove_dropout:
                x = self.dropout(x)
        x = self.convs[-1](x, adj_t)
        return x.log_softmax(dim=-1)


class SAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(SAGE, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(SAGEConv(in_channels, hidden_channels))
        self.bns = torch.nn.ModuleList()
        self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.convs.append(SAGEConv(hidden_channels, out_channels))
        self.dropout = torch.nn.Dropout(p=dropout)
        self.act = torch.nn.ReLU()


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if not config.debug_remove_bn:
            for bn in self.bns:
                bn.reset_parameters()


    def forward(self, x, adj_t):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, adj_t)
            x = self.bns[i](x)
            x = self.act(x)
            if not config.debug_remove_dropout:
                x = self.dropout(x)
        x = self.convs[-1](x, adj_t)
        return x.log_softmax(dim=-1)
mem_speed_bench/arxiv/gnn.py
saintslab/i-Exact
import sys
import os
import numpy as np
from torch.optim import optimizer
sys.path.append(os.getcwd())
import argparse
import random
import time
import yaml
import math
import warnings
from copy import deepcopy

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from torch.cuda.amp import autocast, GradScaler

from torch_geometric.nn.conv.gcn_conv import gcn_norm
from ogb.nodeproppred import Evaluator
from carbontracker.tracker import CarbonTracker

from exact import get_memory_usage, compute_tensor_bytes, exp_recorder, config, QModule

import models
import torch_geometric.transforms as T
from data import get_data
from logger import Logger


EPSILON = 1 - math.log(2)
MB = 1024**2
GB = 1024**3


parser = argparse.ArgumentParser()
parser.add_argument('--conf', type=str, required=True)
parser.add_argument('--root', type=str, default='./data')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--debug_mem', action='store_true',
                    help='whether to debug the memory usage')
parser.add_argument('--n_bits', type=int, default=8)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--simulate', action='store_true')
parser.add_argument('--act_fp', action='store_true')
parser.add_argument('--kept_frac', type=float, default=1.0)
parser.add_argument('--amp', help='whether to enable apx mode', action='store_true')
parser.add_argument('--test_speed', action='store_true', help='whether to test the speed and throughout')
parser.add_argument('--test_co2', action='store_true', help='whether to test the carbon footprint')
parser.add_argument('--col_size', type=int, default=0)
parser.add_argument('--lo', type=float, default=1.0)
parser.add_argument('--use_optimal_lo', action='store_true')

 
def add_labels(feat, labels, n_classes, idx):
    onehot = torch.zeros([feat.shape[0], n_classes], device='cuda')
    onehot[idx, labels[idx, 0]] = 1
    return torch.cat([feat, onehot], dim=-1)


def custom_loss_function(x, labels):
    y = F.cross_entropy(x, labels[:, 0], reduction="none")
    y = torch.log(EPSILON + y) - math.log(EPSILON)
    return torch.mean(y)


def adjust_learning_rate(optimizer, lr, epoch):
    if epoch <= 50:
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr * epoch / 50


def get_optimizer(model_config, model):
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    return optimizer


@torch.no_grad()
def call_model(model, data, n_classes, use_labels=False, amp_mode=False):
    model.eval()
    if use_labels:
        feat = add_labels(data.x, data.y, n_classes, data.train_mask)
    else:
        feat = data.x
    with autocast(enabled=amp_mode):
        model(feat, data.adj_t)


@torch.no_grad()
def test(model, data, evaluator, n_classes, use_labels=False, amp_mode=False):
    model.eval()
    if use_labels:
        feat = add_labels(data.x, data.y, n_classes, data.train_mask)
    else:
        feat = data.x
    with autocast(enabled=amp_mode):
        out = model(feat, data.adj_t)
    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': data.y[data.train_mask],
        'y_pred': y_pred[data.train_mask],
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': data.y[data.val_mask],
        'y_pred': y_pred[data.val_mask],
    })['acc']
    test_acc = evaluator.eval({
        'y_true': data.y[data.test_mask],
        'y_pred': y_pred[data.test_mask],
    })['acc']

    return train_acc, valid_acc, test_acc


def main():
    args = parser.parse_args()
    with open(args.conf, 'r') as fp:
        model_config = yaml.load(fp, Loader=yaml.FullLoader)
    args.model = model_config['name']
    assert args.model.lower() in ['gcn', 'sage', 'gat', 'gcn2']
    if args.simulate:
        print('=' * 30 + 'Simulation' + '=' * 30 )
        config.simulate = True
    if args.act_fp and args.kept_frac == 1.0:
        use_qmodule = False
        print('=' * 30 + 'Do not compress activation' + '=' * 30)
        config.compress_activation = False
    else:
        use_qmodule = True
        config.compress_activation = not args.act_fp
        config.activation_compression_bits = [args.n_bits]
        precision = 'FP32' if not config.compress_activation\
                     else f'INT{config.activation_compression_bits[0]}'
        print('=' * 30 + f'Store activation in {precision}' + '=' * 30)
    print(f'use qmodule: {use_qmodule}')
    if args.kept_frac < 1.0:
        config.kept_frac = args.kept_frac
    print(f'kept frac: {args.kept_frac}')
    if args.amp:
        config.amp = True
        scaler = GradScaler()
    else:
        config.amp = False
        scaler = None
    print(f'amp mode: {config.amp}')
    use_custom_loss = model_config.get('use_custom_loss', True)
    if use_custom_loss:
        print('use the custom loss function')
    else:
        print('use crossentropy as the loss function')
    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        print("Use GPU {} for training".format(args.gpu))

    if args.amp:
        print('activate amp mode')
        config.amp = True
        scaler = GradScaler()
    else:
        scaler = None

    if args.col_size is not None:
        config.col_size = args.col_size

    if args.lo is not None:
        config.lo = args.lo

    if args.use_optimal_lo:
        config.use_optimal_lo = True

    use_labels, n_label_iters, mask_rate = model_config.get('use_labels', False), \
                                           model_config.get('n_label_iters', 0.), \
                                           model_config.get('mask_rate', 0.5)
    grad_norm = model_config.get('grad_norm', None)
    print(f'clipping grad norm: {grad_norm}')
    if not use_labels and n_label_iters > 0:
        raise ValueError("'use_labels' must be enabled when n_label_iters > 0")
    torch.cuda.set_device(args.gpu)

    data, num_features, num_classes = get_data(args.root, 'arxiv')
    data = T.ToSparseTensor()(data)
    num_features = num_features + num_classes if use_labels else num_features
    evaluator = Evaluator(name='ogbn-arxiv')
    logger = Logger(args.runs, args)
    GNN = getattr(models, args.model)
    model = GNN(in_channels=num_features, out_channels=num_classes, **model_config['architecture'])
    if model_config['loop']:
        t = time.perf_counter()
        print('Adding self-loops...', end=' ', flush=True)
        data.adj_t = data.adj_t.set_diag()
        print(f'Done! [{time.perf_counter() - t:.2f}s]')
    
    if model_config['normalize']:
        t = time.perf_counter()
        print('Normalizing data...', end=' ', flush=True)
        data.adj_t = gcn_norm(data.adj_t, add_self_loops=False)
        print(f'Done! [{time.perf_counter() - t:.2f}s]')

    if use_qmodule:
        print('convert the model')
        model = QModule(model)
    print(model)
    model.cuda(args.gpu)
    if args.debug_mem:
        print("========== Model and Optimizer Only ===========")
        optimizer = get_optimizer(model_config, model)
        optimizer.zero_grad()
        usage = get_memory_usage(args.gpu, True)
        exp_recorder.record("network", args.model)
        exp_recorder.record("model_only", usage / MB, 4)
        print("========== Load data to GPU ===========")
        data.adj_t.fill_cache_()
        data = data.to('cuda')
        init_mem = get_memory_usage(args.gpu, True)
        data_mem = init_mem / MB - exp_recorder.val_dict['model_only']
        exp_recorder.record("data", init_mem / MB - exp_recorder.val_dict['model_only'], 4)
        model.reset_parameters()
        model.train()
        with autocast(enabled=args.amp):
            out = model(data.x, data.adj_t)[data.train_idx]
            loss = F.nll_loss(out, data.y.squeeze(1)[data.train_idx])
        print("========== Before Backward ===========")
        before_backward = get_memory_usage(args.gpu, True)
        act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss, out])

        res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                           data_mem,
                                                                           act_mem / MB)
        print(f'max allocated mem (MB): {torch.cuda.max_memory_allocated(0) / MB}')
        print(res)
        if args.amp:
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            optimizer.step()
        del loss, out
        print("========== After Backward ===========")
        after_backward = get_memory_usage(args.gpu, True)
        total_mem = before_backward + (after_backward - init_mem)
        res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (total_mem / MB,
                                                                           data_mem,
                                                                           act_mem / MB)
        print(res)
        print(f'max allocated mem (MB): {torch.cuda.max_memory_allocated(0) / MB}')
        exp_recorder.record("total", total_mem / MB, 2)
        exp_recorder.record("activation", act_mem / MB, 2)
        exp_recorder.dump('mem_results.json')
        torch.cuda.synchronize()
        s_time = time.time()
        if args.test_speed:
            model.reset_parameters()
            optimizer.zero_grad()
            epoch_per_sec = []
            for i in range(100):
                t = time.time()
                optimizer.zero_grad()
                out = model(data.x, data.adj_t)[data.train_idx]
                loss = F.nll_loss(out, data.y.squeeze(1)[data.train_idx])
                loss.backward()
                optimizer.step()
                duration = time.time() - t
                epoch_per_sec.append(duration)
                # print(f'epoch {i}, duration: {duration} sec')
            torch.cuda.synchronize()
            print(f'epoch/s: {100/(time.time() - s_time) }')

        if args.test_co2:
            model.reset_parameters()
            optimizer.zero_grad()
            tracker = CarbonTracker(epochs=100)
            for i in range(100):
                tracker.epoch_start()

                optimizer.zero_grad()
                out = model(data.x, data.adj_t)[data.train_idx]
                loss = F.nll_loss(out, data.y.squeeze(1)[data.train_idx])
                loss.backward()
                optimizer.step()

                tracker.epoch_end()
            tracker.stop()
        exit()

    data = data.to('cuda')
    for run in range(args.runs):
        best_val = 0
        best_model = None

        model.reset_parameters()
        optimizer = get_optimizer(model_config, model)
        for epoch in range(1, 1 + model_config['epochs']):
            # ==== train the model ====
            if model_config.get('adjust_lr', False):
                adjust_learning_rate(optimizer, model_config['lr'], epoch)
            model.train()
            optimizer.zero_grad()
            with autocast(enabled=args.amp):
                if use_labels:
                    mask = torch.rand(data.train_idx.shape) < mask_rate
                    train_labels_idx = data.train_idx[mask]
                    train_pred_idx = data.train_idx[~mask]
                    feat = add_labels(data.x, data.y, num_classes, train_labels_idx)
                else:
                    train_pred_idx = data.train_idx
                    feat = data.x
                out = model(feat, data.adj_t)
                if use_custom_loss:
                    loss = custom_loss_function(out[train_pred_idx], data.y[train_pred_idx])
                else:
                    loss = F.cross_entropy(out[train_pred_idx], data.y.squeeze(1)[train_pred_idx])
            del feat
            torch.cuda.empty_cache()
            if args.amp:
                scaler.scale(loss).backward()
                if grad_norm is not None:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                if grad_norm is not None:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
                optimizer.step()
            # ===========================
            result = test(model, data, evaluator, num_classes, use_labels, args.amp)
            logger.add_result(run, result)
            if  model_config['log_steps'] > 0 and epoch % model_config['log_steps'] == 0:
                train_acc, valid_acc, test_acc = result
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Loss: {loss.item():.4f}, '
                      f'Train: {100 * train_acc:.2f}%, '
                      f'Valid: {100 * valid_acc:.2f}% '
                      f'Test: {100 * test_acc:.2f}%')
            
            if result[1] > best_val:
                #print(f'{epoch}: New best model found')
                best_model = deepcopy(model)
                best_val = result[1]
        
            if epoch == model_config['epochs']:
                config.last_epoch = True
                call_model(best_model, data, num_classes, use_labels, args.amp)
                config.last_epoch = False

        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == '__main__':
    main()
mem_speed_bench/arxiv/train_full_batch.py
saintslab/i-Exact
import sys
import os
import numpy as np
from torch.optim import optimizer
sys.path.append(os.getcwd())
import argparse
import random
import time
import yaml
import math
import warnings

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from torch_geometric.data import GraphSAINTRandomWalkSampler, ClusterData, ClusterLoader
import torch_geometric.transforms as T


from torch_geometric.nn.conv.gcn_conv import gcn_norm
from ogb.nodeproppred import Evaluator

from exact import get_memory_usage, compute_tensor_bytes, exp_recorder, config, QModule

import models
from data import get_data
from logger import Logger


EPSILON = 1 - math.log(2)
MB = 1024**2
GB = 1024**3

SAMPLER = {'graph_saint': GraphSAINTRandomWalkSampler, 'cluster_gcn': ClusterData}

parser = argparse.ArgumentParser()
parser.add_argument('--conf', type=str, required=True)
parser.add_argument('--root', type=str, default='./data')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--debug_mem', action='store_true',
                    help='whether to debug the memory usage')
parser.add_argument('--n_bits', type=int, default=8)
parser.add_argument('--num_workers', type=int, default=12)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--simulate', action='store_true')
parser.add_argument('--act_fp', action='store_true')
parser.add_argument('--kept_frac', type=float, default=1.0)
parser.add_argument('--test_speed', action='store_true', help='whether to test the speed and throughout')


def custom_loss_function(x, labels):
    y = F.cross_entropy(x, labels, reduction="none")
    y = torch.log(EPSILON + y) - math.log(EPSILON)
    return torch.mean(y)


def train(model, loader, optimizer, loss_op, grad_norm):
    model.train()
    total_loss = 0
    for data in loader:
        # s_time = time.time()
        data = T.ToSparseTensor()(data.to('cuda'))
        optimizer.zero_grad()
        out = model(data.x, data.adj_t)
        loss = loss_op(out[data.train_mask], data.y.squeeze(1)[data.train_mask])
        del data
        loss.backward()
        if grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
        optimizer.step()
        total_loss += loss.item()
        # print(f'used time: {time.time() - s_time}')
    return total_loss / len(loader)


def get_optimizer(model_config, model):
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    return optimizer


@torch.no_grad()
def test(model, data, evaluator):
    model.eval()
    feat = data.x
    out = model(feat, data.adj_t)
    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': data.y[data.train_mask],
        'y_pred': y_pred[data.train_mask],
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': data.y[data.val_mask],
        'y_pred': y_pred[data.val_mask],
    })['acc']
    test_acc = evaluator.eval({
        'y_true': data.y[data.test_mask],
        'y_pred': y_pred[data.test_mask],
    })['acc']

    return train_acc, valid_acc, test_acc


def main():
    args = parser.parse_args()
    with open(args.conf, 'r') as fp:
        model_config = yaml.load(fp, Loader=yaml.FullLoader)
    args.model = model_config['name']
    assert args.model.lower() in ['graph_saint', 'cluster_gcn']
    if args.simulate:
        print('=' * 30 + 'Simulation' + '=' * 30 )
        config.simulate = True
    if args.act_fp and args.kept_frac == 1.0:
        use_qmodule = False
        print('=' * 30 + 'Do not compress activation' + '=' * 30)
        config.compress_activation = False
    else:
        use_qmodule = True
        config.compress_activation = not args.act_fp
        config.activation_compression_bits = [args.n_bits]
        precision = 'FP32' if not config.compress_activation\
                     else f'INT{config.activation_compression_bits[0]}'
        print('=' * 30 + f'Store activation in {precision}' + '=' * 30)
    print(f'use qmodule: {use_qmodule}')
    if args.kept_frac < 1.0:
        config.kept_frac = args.kept_frac
    print(f'kept frac: {args.kept_frac}')
    use_custom_loss = model_config.get('use_custom_loss', True)
    if use_custom_loss:
        print('use the custom loss function')
    else:
        print('use crossentropy as the loss function')
    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        print("Use GPU {} for training".format(args.gpu))

    grad_norm = model_config.get('grad_norm', None)
    print(f'clipping grad norm: {grad_norm}')
    torch.cuda.set_device(args.gpu)

    data, num_features, num_classes = get_data(args.root, 'arxiv')
    sampler_data = data
    print('converting data form...')
    s_time = time.time()
    data = T.ToSparseTensor()(data.clone())
    data = data.to('cuda')
    print(f'done. used {time.time() - s_time} sec')

    sampler_cls = SAMPLER[model_config['name']]
    s_time = time.time()
    print("=" * 30 + f'Prepare {sampler_cls.__name__} for training' + '=' * 30)
    if model_config['name'] == 'graph_saint':
        num_steps = int(sampler_data.num_nodes / model_config['sampler']['batch_size'] / model_config['sampler']['walk_length'])
        print(f'num steps: {num_steps}')
        loader = sampler_cls(sampler_data,
                            num_workers=args.num_workers,
                            num_steps=num_steps,
                            **model_config['sampler'])

    elif model_config['name'] == 'cluster_gcn':
        batch_size = model_config['sampler']['batch_size']
        del model_config['sampler']['batch_size']
        cluster_data = sampler_cls(sampler_data, 
                                   **model_config['sampler'])
        loader = ClusterLoader(cluster_data, batch_size=batch_size, 
                               shuffle=True, num_workers=args.num_workers)
    else:
        raise NotImplementedError
    print("=" * 30 + 
          f'Finished Building {sampler_cls.__name__}, used {time.time() - s_time} sec' + 
          '=' * 30)

    evaluator = Evaluator(name='ogbn-arxiv')
    logger = Logger(args.runs, args)
    GNN = getattr(models, model_config['arch_name'])
    model = GNN(in_channels=num_features, out_channels=num_classes, **model_config['architecture'])

    if use_qmodule:
        print('convert the model')
        model = QModule(model)
    print(model)
    model.cuda(args.gpu)
    loss_op = custom_loss_function if use_custom_loss else F.cross_entropy
    if args.debug_mem:
        print("========== Model and Optimizer Only ===========")
        optimizer = get_optimizer(model_config, model)
        optimizer.zero_grad()
        usage = get_memory_usage(args.gpu, True)
        exp_recorder.record("network", args.model)
        exp_recorder.record("model_only", usage / MB, 4)
        print("========== Load data to GPU ===========")
        for _data in loader:     
            optimizer.zero_grad()
            _data = T.ToSparseTensor()(_data.to('cuda'))
            _data.adj_t.fill_cache_()
            init_mem = get_memory_usage(args.gpu, True)
            data_mem = init_mem / MB - exp_recorder.val_dict['model_only']
            exp_recorder.record("data", init_mem / MB - exp_recorder.val_dict['model_only'], 4)
            model.reset_parameters()
            model.train()
            out = model(_data.x, _data.adj_t)[_data.train_mask]
            loss = F.nll_loss(out, _data.y.squeeze(1)[_data.train_mask])
            print("========== Before Backward ===========")
            before_backward = get_memory_usage(args.gpu, True)
            act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss, out])

            res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                            data_mem,
                                                                            act_mem / MB)
            print(res)
            del out
            loss.backward()
            optimizer.step()
            exp_recorder.record("activation", act_mem / MB, 2)
            exp_recorder.dump('mem_results.json')
            torch.cuda.synchronize()
            s_time = time.time()
            if args.test_speed:
                model.reset_parameters()
                optimizer.zero_grad()
                epoch_per_sec = []
                for i in range(100):
                    t = time.time()
                    optimizer.zero_grad()
                    out = model(data.x, data.adj_t)[data.train_idx]
                    loss = F.nll_loss(out, data.y.squeeze(1)[data.train_idx])
                    loss.backward()
                    optimizer.step()
                    duration = time.time() - t
                    epoch_per_sec.append(duration)
                    print(f'epoch {i}, duration: {duration} sec')
                print(f's/epoch: {np.mean(epoch_per_sec)}')
                torch.cuda.synchronize()
                print(f'epoch/s: {100/(time.time() - s_time) }')
            exit()

    data = data.to('cuda')
    for run in range(args.runs):
        model.reset_parameters()
        optimizer = get_optimizer(model_config, model)
        for epoch in range(1, 1 + model_config['epochs']):
            loss = train(model, loader, optimizer, loss_op, grad_norm)
            result = test(model, data, evaluator)
            logger.add_result(run, result)
            if  model_config['log_steps'] > 0 and epoch % model_config['log_steps'] == 0:
                train_acc, valid_acc, test_acc = result
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Loss: {loss:.4f}, '
                      f'Train: {100 * train_acc:.2f}%, '
                      f'Valid: {100 * valid_acc:.2f}% '
                      f'Test: {100 * test_acc:.2f}%')

        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == '__main__':
    main()
mem_speed_bench/arxiv/train_mini_batch.py
saintslab/i-Exact
from typing import Tuple

import torch_geometric.transforms as T
from torch_geometric.data import Data
from torch_geometric.datasets import (Yelp, Flickr, Reddit2, Reddit)
from torch_geometric.utils import to_undirected
from ogb.nodeproppred import PygNodePropPredDataset
from utils import index2mask


def get_arxiv(root: str) -> Tuple[Data, int, int]:
    dataset = PygNodePropPredDataset('ogbn-arxiv', f'{root}/OGB')
    data = dataset[0]
    # data.adj_t = data.adj_t.to_symmetric()
    data.edge_index = to_undirected(data.edge_index)
    data.x = data.x.contiguous()
    data.node_year = None
    split_idx = dataset.get_idx_split()
    data.train_mask = index2mask(split_idx['train'], data.num_nodes)
    data.val_mask = index2mask(split_idx['valid'], data.num_nodes)
    data.test_mask = index2mask(split_idx['test'], data.num_nodes)
    data.train_idx = split_idx['train']
    return data, dataset.num_features, dataset.num_classes


def get_products(root: str) -> Tuple[Data, int, int]:
    dataset = PygNodePropPredDataset('ogbn-products', f'{root}/OGB')
    data = dataset[0]
    data.x = data.x.contiguous()
    split_idx = dataset.get_idx_split()
    data.train_mask = index2mask(split_idx['train'], data.num_nodes)
    data.val_mask = index2mask(split_idx['valid'], data.num_nodes)
    data.test_mask = index2mask(split_idx['test'], data.num_nodes)
    return data, dataset.num_features, dataset.num_classes


def get_yelp(root: str) -> Tuple[Data, int, int]:
    dataset = Yelp(f'{root}/YELP')
    data = dataset[0]
    data.x = (data.x - data.x.mean(dim=0)) / data.x.std(dim=0)
    return data, dataset.num_features, dataset.num_classes


def get_flickr(root: str) -> Tuple[Data, int, int]:
    dataset = Flickr(f'{root}/Flickr')
    return dataset[0], dataset.num_features, dataset.num_classes


def get_reddit(root: str) -> Tuple[Data, int, int]:
    dataset = Reddit(f'{root}/Reddit')
    data = dataset[0]
    return data, dataset.num_features, dataset.num_classes

def get_reddit2(root: str) -> Tuple[Data, int, int]:
    dataset = Reddit2(f'{root}/Reddit2')
    data = dataset[0]
    data.x = (data.x - data.x.mean(dim=0)) / data.x.std(dim=0)
    return data, dataset.num_features, dataset.num_classes

def get_data(root: str, name: str) -> Tuple[Data, int, int]:
    if name.lower() == 'reddit':
        return get_reddit(root)
    elif name.lower() == 'reddit2':
        return get_reddit2(root)
    elif name.lower() == 'flickr':
        return get_flickr(root)
    elif name.lower() == 'yelp':
        return get_yelp(root)
    elif name.lower() == 'arxiv':
        return get_arxiv(root)
    elif name.lower() == 'products':
        return get_products(root)
    else:
        raise NotImplementedError
mem_speed_bench/data.py
saintslab/i-Exact
import torch
import numpy as np
import matplotlib.pyplot as plt
import yaml
import os
import pdb
import itertools
import seaborn as sns

FONTSIZE = 14
LABELSIZE = 15
LEGENDSIZE = 14
TEXTSIZE = 20
MARKERSIZE = 64

markers = 'ov*d'
colors = 'rbgm'

font_config = {'font.size': FONTSIZE, 'font.family': 'sans-serif'}
plt.rcParams.update(font_config)
plt.rcParams["figure.figsize"] = (6.4, 2.8)


fracs = [0.125, 0.25, 0.5, 1.0]
bit = 2
dataset = 'arxiv'
archs = ['gcn', 'sage', 'gcn2']
leg = []
throughout = {}
run_time = {}

for arch in archs:
    throughout[arch], run_time[arch] = [], []
    for frac in fracs:
        precision = f'int{bit}'
        log_path = f'./{dataset}_{arch}_{precision}_frac{frac}.txt'
        lines = open(log_path, 'r').readlines()
        try:
            th = float(lines[-1].split('epoch/s: ')[1][:-2])
            rt = float(lines[-2].split('s/epoch: ')[1][:-2])
            throughout[arch].append(th)
            run_time[arch].append(rt)
        except:
            print(arch, frac)
    bl_path = f'./{dataset}_{arch}_fp32_frac1.0.txt'
    lines = open(bl_path, 'r').readlines()
    th = float(lines[-1].split('epoch/s: ')[1][:-2])
    rt = float(lines[-2].split('s/epoch: ')[1][:-2])
    throughout[arch].append(th)
    run_time[arch].append(rt)
print(throughout)

X = np.arange(5)
seq = ['gcn', 'sage', 'gcn2']

for i in range(1):
    plt.bar(X, throughout[seq[i]])
pdb.set_trace()
plt.savefig(f'test.png', bbox_inches='tight')


# for frac, c, m in zip(fracs, colors, markers):
#     plt.plot(range(len(means[frac])), means[frac], f'-{m}k', color=c, markersize=6,  linewidth=2)
#     plt.xticks(np.arange(3), ['INT2', 'INT4', 'INT8'])
#     # if frac == 0.5:
#     plt.fill_between(range(len(means[frac])), 
#                     np.array(means[frac]) - np.array(stds[frac]), 
#                     np.array(means[frac]) + np.array(stds[frac]), color=c, alpha=0.1)

# plt.plot(range(3), [baselines[dataset]] * 3, color='k')
# leg = [f'D/R = {int(1/frac)}' for frac in fracs] + ['baseline']
# plt.legend(leg, prop={'size': 12}, bbox_to_anchor=(1, 1))
# # if arch == 'sage':
# #     plt.legend(leg, prop={'size': 12}, bbox_to_anchor=(1, 1))
# # if arch == 'sage':
# #     plt.yticks(np.arange(71.0, 72.1, 0.2))
# # elif arch == 'gcn':
# #     plt.yticks(np.arange(71.2, 72.2, 0.2))
# # else:
# #     raise NotImplemented
# filename = f'./{dataset}_{arch}_results_norm{norm}'
# if large_batch:
#     filename += '_large_batch'
# plt.savefig(f'{filename}.png', bbox_inches='tight')

mem_speed_bench/mem_speed_logs_lat_version/gather_running_speed.py
saintslab/i-Exact
from .gcn import GCN
from .sage import SAGE
from .gat import GAT
from .kgnn import KGNN
from .gcn2 import GCN2

__all__ = [
    'GCN',
    'SAGE',
    'GAT',
    'KGNN',
    'GCN2'
]
mem_speed_bench/models/__init__.py
saintslab/i-Exact
from typing import Optional
from exact.gatconv import CustomGATConv

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import ModuleList, Linear, BatchNorm1d
from torch_sparse import SparseTensor



class GAT(torch.nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int,
                 out_channels: int, num_layers: int, num_heads: int, 
                 dropout: float = 0.0, input_drop: float = 0.0, 
                 attn_drop: float = 0.0, edge_drop: float = 0.0,
                 use_attn_dst: bool = True, residual: bool = False,
                 batch_norm: bool = False,):
        super(GAT, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dropout = torch.nn.Dropout(p=dropout)
        self.input_dropout = torch.nn.Dropout(p=input_drop)
        self.activation = torch.nn.ReLU()
        self.batch_norm = batch_norm
        self.residual = residual
        self.use_attn_dst = use_attn_dst
        self.num_heads = num_heads
        self.attn_drop = attn_drop
        self.edge_drop = edge_drop
        self.num_layers = num_layers

        self.convs = ModuleList()
        for i in range(num_layers):
            in_dim = hidden_channels * num_heads
            out_dim = hidden_channels
            output_heads = num_heads
            if i == 0:
                in_dim = in_channels
            if i == num_layers - 1:
                out_dim = out_channels
                output_heads = 1
            conv = CustomGATConv(in_dim, out_dim, output_heads, dropout=attn_drop, 
                                 add_self_loops=False, residual=residual, use_attn_dst=use_attn_dst)
            self.convs.append(conv)

        if self.batch_norm:
            self.bns = ModuleList()
            for i in range(num_layers - 1):
                bn = BatchNorm1d(hidden_channels * num_heads)
                self.bns.append(bn)


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if self.batch_norm:
            for bn in self.bns:
                bn.reset_parameters()


    def forward(self, x: Tensor, adj_t: SparseTensor, *args) -> Tensor:
        x = self.input_dropout(x)
        if self.training and self.edge_drop > 0:
            nnz = adj_t.nnz()
            bound = int(nnz * self.edge_drop)
            perm = torch.randperm(nnz).cuda()
            eids = perm[bound:]
            row, col, value = adj_t.storage._row, adj_t.storage._col, adj_t.storage._value
            row, col, value = row[eids], col[eids], value[eids]
            adj_t = SparseTensor(row=row, col=col, value=value, 
                                    sparse_sizes=(adj_t.size(0), adj_t.size(1)), is_sorted=False)
        for idx, conv in enumerate(self.convs[:-1]):
            x = conv(x, adj_t)
            if self.batch_norm:
                x = self.bns[idx](x)
            x = self.activation(x)
            x = self.dropout(x)
        x = self.convs[-1](x, adj_t)
        return x


    @torch.no_grad()
    def forward_layer(self, layer, x, adj_t):
        x = self.convs[layer](x, adj_t)
        if layer < self.num_layers - 1:
            if self.batch_norm:
                x = self.bns[layer](x)
            x = self.activation(x)
        return x
mem_speed_bench/models/gat.py
saintslab/i-Exact
from tqdm import tqdm
import torch
from torch import Tensor
from torch.nn import ModuleList, BatchNorm1d
from torch_sparse import SparseTensor
from torch_geometric.nn import GCNConv



class GCN(torch.nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int,
                 out_channels: int, num_layers: int, dropout: float = 0.0,
                 drop_input: bool = False, batch_norm: bool = False, residual: bool = False):
        super(GCN, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dropout = torch.nn.Dropout(p=dropout)
        self.drop_input = drop_input
        if drop_input:
            self.input_dropout = torch.nn.Dropout(p=dropout)
        self.activation = torch.nn.ReLU()
        self.batch_norm = batch_norm
        self.residual = residual
        self.num_layers = num_layers
        self.convs = ModuleList()
        for i in range(num_layers):
            in_dim = out_dim = hidden_channels
            if i == 0:
                in_dim = in_channels
            if i == num_layers - 1:
                out_dim = out_channels
            conv = GCNConv(in_dim, out_dim, normalize=False)
            self.convs.append(conv)

        if self.batch_norm:
            self.bns = ModuleList()
            for i in range(num_layers - 1):
                bn = BatchNorm1d(hidden_channels)
                self.bns.append(bn)


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if self.batch_norm:
            for bn in self.bns:
                bn.reset_parameters()


    def forward(self, x: Tensor, adj_t: SparseTensor, *args) -> Tensor:
        if self.drop_input:
            x = self.input_dropout(x)
        for idx, conv in enumerate(self.convs[:-1]):
            h = conv(x, adj_t)
            if self.batch_norm:
                h = self.bns[idx](h)
            if self.residual and h.size(-1) == x.size(-1):
                h += x[:h.size(0)]
            x = self.activation(h)
            x = self.dropout(x)

        x = self.convs[-1](x, adj_t)
        return x


    @torch.no_grad()
    def forward_layer(self, layer, x, adj_t):
        if layer != 0:
            if self.drop_input:
                x = self.input_dropout(x)
            x = self.dropout(x)
        h = self.convs[layer](x, adj_t)
        if layer < self.num_layers - 1:
            if self.batch_norm:
                h = self.bns[layer](h)
            if self.residual and h.size(-1) == x.size(-1):
                h += x[:h.size(0)]
            h = self.activation(h)
        return h


    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        pbar = tqdm(total=x_all.size(0) * len(self.convs))
        pbar.set_description('Evaluating')
        for i in range(len(self.convs)):
            xs = []
            for batch_size, n_id, adj in loader:
                edge_index, _, size = adj.to('cuda')
                x = x_all[n_id].to('cuda')
                xs.append(self.forward_layer(i, x, edge_index).cpu())
                pbar.update(batch_size)
            x_all = torch.cat(xs, dim=0)
        pbar.close()
        return x_all    
mem_speed_bench/models/gcn.py
saintslab/i-Exact
from tqdm import tqdm
import torch
import torch.nn.functional as F

from torch import Tensor
from torch.nn import ModuleList, BatchNorm1d, Linear
from torch_sparse import SparseTensor
from torch_geometric.nn import GCN2Conv



class GCN2(torch.nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int,
                 out_channels: int, num_layers: int, alpha: float, theta: float = None, 
                 shared_weights: bool = True, dropout: float = 0.0, drop_input: bool = True,
                 batch_norm: bool = False, residual: bool = False):
        super(GCN2, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dropout = torch.nn.Dropout(p=dropout)
        self.drop_input = drop_input
        if drop_input:
            self.input_dropout = torch.nn.Dropout(p=dropout)
        self.activation = torch.nn.ReLU()
        self.batch_norm = batch_norm
        self.residual = residual
        self.num_layers = num_layers
        self.alpha, self.theta = alpha, theta

        self.lins = ModuleList()
        self.lins.append(Linear(in_channels, hidden_channels))
        self.lins.append(Linear(hidden_channels, out_channels))

        self.convs = ModuleList()
        for i in range(num_layers):
            if theta is None:
                conv = GCN2Conv(hidden_channels, alpha=alpha, theta=None,
                                layer=None, shared_weights=shared_weights,
                                normalize=False, add_self_loops=False)
            else:
                conv = GCN2Conv(hidden_channels, alpha=alpha, theta=theta,
                                layer=i+1, shared_weights=shared_weights,
                                normalize=False, add_self_loops=False)
            self.convs.append(conv)

        if self.batch_norm:
            self.bns = ModuleList()
            for i in range(num_layers):
                bn = BatchNorm1d(hidden_channels)
                self.bns.append(bn)


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if self.batch_norm:
            for bn in self.bns:
                bn.reset_parameters()
        for lin in self.lins:
            lin.reset_parameters()


    def forward(self, x: Tensor, adj_t: SparseTensor, *args) -> Tensor:
        if self.drop_input:
            x = self.input_dropout(x)
        x = x0 = self.activation(self.lins[0](x))
        x = self.dropout(x)
        for idx, conv in enumerate(self.convs[:-1]):
            h = conv(x, x0, adj_t)
            if self.batch_norm:
                h = self.bns[idx](h)
            if self.residual:
                h += x[:h.size(0)]
            x = self.activation(h)
            x = self.dropout(x)

        h = self.convs[-1](x, x0, adj_t)
        if self.batch_norm:
            h = self.bns[-1](h)
        if self.residual:
            h += x[:h.size(0)]
        x = self.activation(h)
        x = self.dropout(x)
        x = self.lins[1](x)
        return x


    @torch.no_grad()
    def forward_layer(self, layer, x, adj_t, state):
        if layer == 0:
            if self.drop_input:
                x = self.input_dropout(x)
            x = x_0 = self.activation(self.lins[0](x))
            state['x_0'] = x_0[:adj_t.size(0)]
        x = self.dropout(x)
        h = self.convs[layer](x, state['x_0'], adj_t)
        if self.batch_norm:
            h = self.bns[layer](h)
        if self.residual and h.size(-1) == x.size(-1):
            h += x[:h.size(0)]
        x = self.activation(h)
        if layer == self.num_layers - 1:
            x = self.dropout(x)
            x = self.lins[1](x)  
        return h


    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        pbar = tqdm(total=x_all.size(0) * len(self.convs))
        pbar.set_description('Evaluating')
        state = {}
        for i in range(len(self.convs)):
            xs = []
            for batch_size, n_id, adj in loader:
                edge_index, _, size = adj.to('cuda')
                x = x_all[n_id].to('cuda')
                xs.append(self.forward_layer(i, x, edge_index, state).cpu())
                pbar.update(batch_size)
            x_all = torch.cat(xs, dim=0)
        pbar.close()
        return x_all    
mem_speed_bench/models/gcn2.py
saintslab/i-Exact
from tqdm import tqdm
import torch
from torch import Tensor
from torch.nn import ModuleList, BatchNorm1d
from torch_sparse import SparseTensor
from torch_geometric.nn import GraphConv



class KGNN(torch.nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int,
                 out_channels: int, num_layers: int, dropout: float = 0.0,
                 batch_norm: bool = False, residual: bool = False, use_linear=False):
        super(KGNN, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dropout = torch.nn.Dropout(p=dropout)
        self.activation = torch.nn.ReLU()
        self.batch_norm = batch_norm
        self.residual = residual
        self.num_layers = num_layers
        self.convs = ModuleList()
        self.use_linear = use_linear
        for i in range(num_layers):
            in_dim = out_dim = hidden_channels
            if i == 0:
                in_dim = in_channels
            if (i == num_layers - 1) and (not use_linear):
                out_dim = out_channels
            conv = GraphConv(in_dim, out_dim)
            self.convs.append(conv)
        if self.batch_norm:
            self.bns = ModuleList()
            for i in range(num_layers - 1):
                bn = BatchNorm1d(hidden_channels)
                self.bns.append(bn)
        if use_linear:
            self.lins = torch.nn.Linear(num_layers * hidden_channels, out_channels)


    def set_aggr(self, aggr):
        for conv in self.convs:
            conv.aggr = aggr


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if self.batch_norm:
            for bn in self.bns:
                bn.reset_parameters()
        if self.use_linear:
            self.lins.reset_parameters()


    def forward(self, x: Tensor, adj_t: SparseTensor, edge_weight=None) -> Tensor:
        if self.use_linear:
            xs = [x]
        for idx, conv in enumerate(self.convs[:-1]):
            h = conv(x, adj_t, edge_weight)
            if self.use_linear:
                xs.append(h)
            if self.batch_norm:
                h = self.bns[idx](h)
            if self.residual and h.size(-1) == x.size(-1):
                h += x[:h.size(0)]
            x = self.activation(h)
            x = self.dropout(x)
        x = self.convs[-1](x, adj_t, edge_weight)
        if self.use_linear:
            xs.append(x)
            x = torch.cat(xs, dim=-1)
            x = self.lin(x)
        return x


    @torch.no_grad()
    def forward_layer(self, layer, x, adj_t):
        if self.use_linear:
            raise NotImplementedError
        if layer != 0:
            x = self.dropout(x)

        h = self.convs[layer](x, adj_t)
        if layer < self.num_layers - 1:
            if self.batch_norm:
                h = self.bns[layer](h)
            if self.residual and h.size(-1) == x.size(-1):
                h += x[:h.size(0)]
            h = self.activation(h)
        return h


    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        pbar = tqdm(total=x_all.size(0) * len(self.convs))
        pbar.set_description('Evaluating')
        for i in range(len(self.convs)):
            xs = []
            for batch_size, n_id, adj in loader:
                edge_index, _, size = adj.to('cuda')
                x = x_all[n_id].to('cuda')
                xs.append(self.forward_layer(i, x, edge_index).cpu())
                pbar.update(batch_size)
            x_all = torch.cat(xs, dim=0)
        pbar.close()
        return x_all    
mem_speed_bench/models/kgnn.py
saintslab/i-Exact
from re import I
from tqdm import tqdm

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import ModuleList, Linear, BatchNorm1d
from torch_sparse import SparseTensor
from torch_geometric.nn import SAGEConv



class SAGE(torch.nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int,
                 out_channels: int, num_layers: int, dropout: float = 0.0,
                 batch_norm: bool = False, residual: bool = False, use_linear=False):
        super(SAGE, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dropout = torch.nn.Dropout(p=dropout)
        self.activation = torch.nn.ReLU()
        self.batch_norm = batch_norm
        self.residual = residual
        self.num_layers = num_layers
        self.convs = ModuleList()
        self.use_linear = use_linear
        if self.use_linear:
            self.lins = torch.nn.ModuleList()
        for i in range(num_layers):
            in_dim = out_dim = hidden_channels
            if i == 0:
                in_dim = in_channels
            if i == num_layers - 1:
                out_dim = out_channels
            conv = SAGEConv(in_dim, out_dim)
            self.convs.append(conv)
            if self.use_linear:
                self.lins.append(torch.nn.Linear(in_dim, out_dim, bias=False))

        if self.batch_norm:
            self.bns = ModuleList()
            for i in range(num_layers - 1):
                bn = BatchNorm1d(hidden_channels)
                self.bns.append(bn)


    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if self.batch_norm:
            for bn in self.bns:
                bn.reset_parameters()
        if self.use_linear:
            for lin in self.lins:
                lin.reset_parameters()


    def forward(self, x: Tensor, adj_t: SparseTensor, *args) -> Tensor:
        for idx, conv in enumerate(self.convs[:-1]):
            h = conv(x, adj_t)
            if self.use_linear:
                linear = self.lins[idx](x)
                h = h + linear
            if self.batch_norm:
                h = self.bns[idx](h)
            if self.residual and h.size(-1) == x.size(-1):
                h += x[:h.size(0)]
            x = self.activation(h)
            x = self.dropout(x)
        h = self.convs[-1](x, adj_t)
        if self.use_linear:
            linear = self.lins[-1](x)
            x = h + linear
        else:
            x = h
        return x


    @torch.no_grad()
    def forward_layer(self, layer, x, adj_t, size):
        if self.use_linear:
            raise NotImplementedError
        if layer != 0:
            x = self.dropout(x)
        x_target = x[:size[1]]
        h = self.convs[layer]((x, x_target), adj_t)
        if layer < self.num_layers - 1:
            if self.batch_norm:
                h = self.bns[layer](h)
            if self.residual and h.size(-1) == x.size(-1):
                h += x[:h.size(0)]
            h = F.relu(h)
        return h
    

    @torch.no_grad()
    def mini_inference(self, x_all, loader):
        pbar = tqdm(total=x_all.size(0) * len(self.convs))
        pbar.set_description('Evaluating')
        for i in range(len(self.convs)):
            xs = []
            for batch_size, n_id, adj in loader:
                edge_index, _, size = adj.to('cuda')
                x = x_all[n_id].to('cuda')
                xs.append(self.forward_layer(i, x, edge_index, size).cpu())
                pbar.update(batch_size)
            x_all = torch.cat(xs, dim=0)
        pbar.close()
        return x_all    
mem_speed_bench/models/sage.py
saintslab/i-Exact
import sys
import os
import numpy as np
from torch.autograd import grad
sys.path.append(os.getcwd())
import argparse
import random
import time
import warnings
import yaml
import pdb
from copy import deepcopy

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn

from torch_geometric.utils import subgraph
from torch_geometric.utils import degree
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from utils import cast_int32
from carbontracker.tracker import CarbonTracker
from exact import get_memory_usage, compute_tensor_bytes, exp_recorder, config, QModule
import models
from data import get_data
from logger import Logger
from sklearn.metrics import f1_score
import torch_geometric.transforms as T

MB = 1024**2
GB = 1024**3


parser = argparse.ArgumentParser()
parser.add_argument('--conf', type=str, required=True, 
                    help='the path to the configuration file')
parser.add_argument('--dataset', type=str, required=True, 
                    help='the name of the applied dataset')
parser.add_argument('--root', type=str, default='./data')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--num_workers', type=int, default=12)
parser.add_argument('--n_bits', type=int, default=8)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--act_fp', action='store_true')
parser.add_argument('--kept_frac', type=float, default=1.0)
parser.add_argument('--grad_norm', type=float, default=None)
parser.add_argument('--inductive', action='store_true')
parser.add_argument('--debug_mem', action='store_true')
parser.add_argument('--test_speed', action='store_true')
parser.add_argument('--simulate', action='store_true')
parser.add_argument('--test_co2', action='store_true', help='whether to test the carbon footprint')
parser.add_argument('--col_size', type=int, default=0)
parser.add_argument('--lo', type=float, default=1.0)
parser.add_argument('--use_optimal_lo', action='store_true')


def get_optimizer(model_config, model):
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    return optimizer


def to_inductive(data):
    mask = data.train_mask
    data.x = data.x[mask]
    data.y = data.y[mask]
    data.train_mask = data.train_mask[mask]
    data.test_mask = None
    data.edge_index, _ = subgraph(mask, data.edge_index, None,
                                  relabel_nodes=True, num_nodes=data.num_nodes)
    data.num_nodes = mask.sum().item()
    return data


def train(model, optimizer, data, loss_op, grad_norm):
    model.train()
    # s_time = time.time()
    optimizer.zero_grad()
    out = model(data.x, data.adj_t)
    loss = loss_op(out[data.train_mask], data.y[data.train_mask])
    del data
    loss.backward()
    if grad_norm is not None:
        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
    optimizer.step()
    # print(f'used time: {time.time() - s_time}')
    return loss.item()


def compute_micro_f1(logits, y, mask=None) -> float:
    if mask is not None:
        logits, y = logits[mask], y[mask]

    if y.dim() == 1:
        return int(logits.argmax(dim=-1).eq(y).sum()) / y.size(0)
        
    else:
        y_pred = logits > 0
        y_true = y > 0.5

        tp = int((y_true & y_pred).sum())
        fp = int((~y_true & y_pred).sum())
        fn = int((y_true & ~y_pred).sum())

        try:
            precision = tp / (tp + fp)
            recall = tp / (tp + fn)
            return 2 * (precision * recall) / (precision + recall)
        except ZeroDivisionError:
            return 0.

@torch.no_grad()
def call_model(model, data):
    model.eval()
    model(data.x, data.adj_t)

@torch.no_grad()
def test(model, data):
    model.eval()
    out = model(data.x, data.adj_t)
    y_true = data.y
    train_acc = compute_micro_f1(out, y_true, data.train_mask)
    valid_acc = compute_micro_f1(out, y_true, data.val_mask)
    test_acc = compute_micro_f1(out, y_true, data.test_mask)
    return train_acc, valid_acc, test_acc


def main():
    global args 
    args = parser.parse_args()
    with open(args.conf, 'r') as fp:
        model_config = yaml.load(fp, Loader=yaml.FullLoader)
        name = model_config['name']
        loop = model_config.get('loop', False)
        normalize = model_config.get('norm', False)
        if args.dataset == 'reddit2':
            model_config = model_config['params']['reddit']
        else:
            model_config = model_config['params'][args.dataset]
        model_config['name'] = name
        model_config['loop'] = loop
        model_config['normalize'] = normalize

    print(f'model config: {model_config}')
    if args.simulate:
        print('=' * 30 + 'Simulation' + '=' * 30 )
        config.simulate = True
    # if args.dataset in ['yelp', 'pubmed']:
    if args.dataset == 'yelp':
        multi_label = True
    else:
        multi_label = False
    print(f'clipping grad norm: {args.grad_norm}')
    args.model = model_config['arch_name']
    assert model_config['name'] in ['GCN', 'SAGE', 'GCN2']
    if args.act_fp and args.kept_frac == 1.0:
        use_qmodule = False
        print('=' * 30 + 'Do not compress activation' + '=' * 30)
        config.compress_activation = False
    else:
        use_qmodule = True
        config.compress_activation = not args.act_fp
        config.activation_compression_bits = [args.n_bits]
        precision = 'FP32' if not config.compress_activation\
                     else f'INT{config.activation_compression_bits[0]}'
        print('=' * 30 + f'Store activation in {precision}' + '=' * 30)
    print(f'use qmodule: {use_qmodule}')
    if args.kept_frac < 1.0:
        config.kept_frac = args.kept_frac
    print(f'kept frac: {args.kept_frac}')

    if args.col_size is not None:
        config.col_size = args.col_size

    if args.perc is not None:
        config.perc = args.perc

    if args.lo is not None:
        config.lo = args.lo

    if args.use_optimal_lo:
        config.use_optimal_lo = True

    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        print("Use GPU {} for training".format(args.gpu))

    torch.cuda.set_device(args.gpu)
    data, num_features, num_classes = get_data(args.root, args.dataset)

    GNN = getattr(models, model_config['arch_name'])
    model = GNN(in_channels=num_features, out_channels=num_classes, **model_config['architecture'])
    if use_qmodule:
        print('convert the model')
        model = QModule(model)
    loss_op = F.binary_cross_entropy_with_logits if multi_label else F.cross_entropy
    print(model)
    model.cuda(args.gpu)

    if args.debug_mem:
        print("========== Model and Optimizer only ===========")
        optimizer = get_optimizer(model_config, model)
        optimizer.zero_grad()
        model.reset_parameters()
        model.train()
        usage = get_memory_usage(args.gpu, False)
        exp_recorder.record("network", args.model)
        exp_recorder.record("model_only", usage / MB, 4)
        print("========== Load data to GPU ===========")
        print('converting data form...')
        s_time = time.time()
        data = T.ToSparseTensor()(data.to('cuda'))
        print(f'done. used {time.time() - s_time} sec')

        if model_config['loop']:
            t = time.perf_counter()
            print('Adding self-loops...', end=' ', flush=True)
            data.adj_t = data.adj_t.set_diag()
            print(f'Done! [{time.perf_counter() - t:.2f}s]')
        
        if model_config['normalize']:
            t = time.perf_counter()
            print('Normalizing data...', end=' ', flush=True)
            data.adj_t = gcn_norm(data.adj_t, add_self_loops=False)
            print(f'Done! [{time.perf_counter() - t:.2f}s]')

        if args.inductive:
            print('inductive learning mode')
            data = to_inductive(data)
        print(f'arch: %s' % model_config['arch_name'])
        # if model_config['arch_name'] in ['GCN', 'GCN2']:
        #     data.adj_t.storage.rowptr()
        #     data.adj_t.storage.col()
        #     data.adj_t.storage.colptr()
        #     data.adj_t.storage.csr2csc()
        # elif model_config['arch_name'] == 'SAGE':
        #     data.adj_t.storage.rowptr()
        #     data.adj_t.storage.col()
        #     data.adj_t.storage.row()
        #     data.adj_t.storage.rowcount()   
        # else:
        #     raise NotImplementedError
        data.adj_t.fill_cache_()
        init_mem = get_memory_usage(args.gpu, False)
        data_mem = init_mem / MB - exp_recorder.val_dict['model_only']
        exp_recorder.record("data", init_mem / MB - exp_recorder.val_dict['model_only'], 4)
        torch.cuda.empty_cache()
        out = model(data.x, data.adj_t)[data.train_mask]
        loss = loss_op(out, data.y[data.train_mask])
        print("========== Before Backward ===========")
        before_backward = get_memory_usage(args.gpu, True)
        act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss, out])

        res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                           data_mem,
                                                                           act_mem / MB)
        print(res)
        loss.backward()
        optimizer.step()
        #print(f'after backward, max allocated mem (MB): %.3f' % ((bwd_peak_compute+before_backward)/MB))
        del loss, out
        print("========== After Backward ===========")
        after_backward = get_memory_usage(args.gpu, True)
        total_mem = before_backward + (after_backward - init_mem)
        res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (total_mem / MB,
                                                                           data_mem,
                                                                           act_mem / MB)
        print(res)
        exp_recorder.record("total", total_mem / MB, 2)
        exp_recorder.record("activation", act_mem / MB, 2)
        # exp_recorder.dump('mem_results.json')
        torch.cuda.synchronize()
        s_time = time.time()
        if args.test_speed:
            model.reset_parameters()
            optimizer.zero_grad()
            epoch_per_sec = []
            for i in range(100):
                t = time.time()
                optimizer.zero_grad()
                out = model(data.x, data.adj_t)[data.train_mask]
                loss = loss_op(out, data.y[data.train_mask])
                loss.backward()
                optimizer.step()
                duration = time.time() - t
                epoch_per_sec.append(duration)
                #print(f'epoch {i}, duration: {duration} sec')
            print(f's/epoch: {np.mean(epoch_per_sec)}')
            torch.cuda.synchronize()
            print(f'epoch/s: {100/(time.time() - s_time) }')

        if args.test_co2:
            model.reset_parameters()
            optimizer.zero_grad()
            tracker = CarbonTracker(epochs=100)
            for i in range(100):
                tracker.epoch_start()

                optimizer.zero_grad()
                out = model(data.x, data.adj_t)[data.train_idx]
                loss = F.nll_loss(out, data.y.squeeze(1)[data.train_idx])
                loss.backward()
                optimizer.step()

                tracker.epoch_end()
            tracker.stop()
        exit()

    print('converting data form...')
    s_time = time.time()
    data = T.ToSparseTensor()(data.to('cuda'))
    print(f'done. used {time.time() - s_time} sec')

    if model_config['loop']:
        t = time.perf_counter()
        print('Adding self-loops...', end=' ', flush=True)
        data.adj_t = data.adj_t.set_diag()
        print(f'Done! [{time.perf_counter() - t:.2f}s]')
    
    if model_config['normalize']:
        t = time.perf_counter()
        print('Normalizing data...', end=' ', flush=True)
        data.adj_t = gcn_norm(data.adj_t, add_self_loops=False)
        print(f'Done! [{time.perf_counter() - t:.2f}s]')

    if args.inductive:
        print('inductive learning mode')
        data = to_inductive(data)
    logger = Logger(args.runs, args)
    for run in range(args.runs):
        best_model = None
        best_val = 0

        model.reset_parameters()
        optimizer = get_optimizer(model_config, model)
        for epoch in range(1, 1 + model_config['epochs']):

            loss = train(model, optimizer, data, loss_op, args.grad_norm)
            # print(f'Run: {run + 1:02d}, '
            #         f'Epoch: {epoch:02d}, '
            #         f'Train Loss: {loss:.4f}')
    
            result = test(model, data)
            logger.add_result(run, result)
            train_acc, valid_acc, test_acc = result
            # print(f'Run: {run + 1:02d}, '
            #         f'Epoch: {epoch:02d}, '
            #         f'Train f1: {100 * train_acc:.2f}%, '
            #         f'Valid f1: {100 * valid_acc:.2f}% '
            #         f'Test f1: {100 * test_acc:.2f}%')
            
            if result[1] > best_val:
                #print(f'{epoch}: New best model found')
                best_model = deepcopy(model)
                best_val = result[1]
        
            if epoch == model_config['epochs']:
                config.last_epoch = True
                call_model(best_model, data)
                config.last_epoch = False

        logger.add_result(run, result)
        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == '__main__':
    main()
mem_speed_bench/non_ogb_datasets/train_full_batch.py
saintslab/i-Exact
import sys
import os

from torch.autograd import grad
sys.path.append(os.getcwd())
import argparse
import random
import time
import warnings
from torch_geometric.data.cluster import ClusterLoader
import yaml
import pdb

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from torch.cuda.amp import autocast, GradScaler

from torch_geometric.utils import subgraph
from torch_geometric.data import GraphSAINTRandomWalkSampler, NeighborSampler, ClusterData
from torch_geometric.utils import degree

from exact import get_memory_usage, compute_tensor_bytes, exp_recorder, config, QModule
import models
from data import get_data
from logger import Logger
from sklearn.metrics import f1_score
import torch_geometric.transforms as T

MB = 1024**2
GB = 1024**3

SAMPLER = {'graph_saint': GraphSAINTRandomWalkSampler, 'cluster_gcn': ClusterData}

parser = argparse.ArgumentParser()
parser.add_argument('--conf', type=str, required=True, 
                    help='the path to the configuration file')
parser.add_argument('--dataset', type=str, required=True, 
                    help='the name of the applied dataset')
parser.add_argument('--root', type=str, default='./data')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--num_workers', type=int, default=12)
parser.add_argument('--n_bits', type=int, default=8)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--act_fp', action='store_true')
parser.add_argument('--kept_frac', type=float, default=1.0)
parser.add_argument('--grad_norm', type=float, default=None)
parser.add_argument('--inductive', action='store_true')
parser.add_argument('--test_speed', action='store_true')
parser.add_argument('--debug_mem', action='store_true')



def get_optimizer(model_config, model):
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    return optimizer


def to_inductive(data):
    mask = data.train_mask
    data.x = data.x[mask]
    data.y = data.y[mask]
    data.train_mask = data.train_mask[mask]
    data.test_mask = None
    data.edge_index, _ = subgraph(mask, data.edge_index, None,
                                  relabel_nodes=True, num_nodes=data.num_nodes)
    data.num_nodes = mask.sum().item()
    return data


def train(model, loader, optimizer, loss_op, grad_norm):
    model.train()
    total_loss = 0
    for data in loader:
        # s_time = time.time()
        data = T.ToSparseTensor()(data.to('cuda'))
        optimizer.zero_grad()
        out = model(data.x, data.adj_t)
        loss = loss_op(out[data.train_mask], data.y[data.train_mask])
        del data
        loss.backward()
        if grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
        optimizer.step()
        total_loss += loss.item()
        # print(f'used time: {time.time() - s_time}')
    return total_loss / len(loader)


def compute_micro_f1(logits, y, mask=None) -> float:
    if mask is not None:
        logits, y = logits[mask], y[mask]

    if y.dim() == 1:
        return int(logits.argmax(dim=-1).eq(y).sum()) / y.size(0)
        
    else:
        y_pred = logits > 0
        y_true = y > 0.5

        tp = int((y_true & y_pred).sum())
        fp = int((~y_true & y_pred).sum())
        fn = int((y_true & ~y_pred).sum())

        try:
            precision = tp / (tp + fp)
            recall = tp / (tp + fn)
            return 2 * (precision * recall) / (precision + recall)
        except ZeroDivisionError:
            return 0.

@torch.no_grad()
def test(model, data):
    model.eval()
    out = model(data.x, data.adj_t)
    y_true = data.y
    train_acc = compute_micro_f1(out, y_true, data.train_mask)
    valid_acc = compute_micro_f1(out, y_true, data.val_mask)
    test_acc = compute_micro_f1(out, y_true, data.test_mask)
    return train_acc, valid_acc, test_acc


def main():
    global args 
    args = parser.parse_args()
    with open(args.conf, 'r') as fp:
        model_config = yaml.load(fp, Loader=yaml.FullLoader)
        name = model_config['name']
        if args.dataset == 'reddit2':
            model_config = model_config['params']['reddit']
        else:
            model_config = model_config['params'][args.dataset]
        model_config['name'] = name
    print(f'model config: {model_config}')
    if args.dataset == 'yelp':
        multi_label = True
    else:
        multi_label = False
    print(f'clipping grad norm: {args.grad_norm}')
    args.model = model_config['arch_name']
    assert model_config['name'] in ['graph_saint', 'cluster_gcn']
    if args.act_fp and args.kept_frac == 1.0:
        use_qmodule = False
        print('=' * 30 + 'Do not compress activation' + '=' * 30)
        config.compress_activation = False
    else:
        use_qmodule = True
        config.compress_activation = not args.act_fp
        config.activation_compression_bits = [args.n_bits]
        precision = 'FP32' if not config.compress_activation\
                     else f'INT{config.activation_compression_bits[0]}'
        print('=' * 30 + f'Store activation in {precision}' + '=' * 30)
    print(f'use qmodule: {use_qmodule}')
    if args.kept_frac < 1.0:
        config.kept_frac = args.kept_frac
    print(f'kept frac: {args.kept_frac}')

    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        print("Use GPU {} for training".format(args.gpu))

    torch.cuda.set_device(args.gpu)
    data, num_features, num_classes = get_data(args.root, args.dataset)
    sampler_data = data
    print('converting data form...')
    s_time = time.time()
    data = T.ToSparseTensor()(data.clone())
    data = data.to('cuda')
    print(f'done. used {time.time() - s_time} sec')

    if args.inductive:
        print('inductive learning mode')
        sampler_data = to_inductive(sampler_data)

    sampler_cls = SAMPLER[model_config['name']]
    s_time = time.time()
    print("=" * 30 + f'Prepare {sampler_cls.__name__} for training' + '=' * 30)
    if model_config['name'] == 'graph_saint':
        num_steps = int(sampler_data.num_nodes / model_config['sampler']['batch_size'] / model_config['sampler']['walk_length'])
        print(f'num steps: {num_steps}')
        loader = sampler_cls(sampler_data,
                            num_workers=args.num_workers,
                            num_steps=num_steps,
                            **model_config['sampler'])

    elif model_config['name'] == 'cluster_gcn':
        batch_size = model_config['sampler']['batch_size']
        del model_config['sampler']['batch_size']
        cluster_data = sampler_cls(sampler_data, 
                                   **model_config['sampler'])
        loader = ClusterLoader(cluster_data, batch_size=batch_size, 
                               shuffle=True, num_workers=args.num_workers)
    else:
        raise NotImplementedError
    print("=" * 30 + 
          f'Finished Building {sampler_cls.__name__}, used {time.time() - s_time} sec' + 
          '=' * 30)

    GNN = getattr(models, model_config['arch_name'])
    model = GNN(in_channels=num_features, out_channels=num_classes, **model_config['architecture'])
    if use_qmodule:
        print('convert the model')
        model = QModule(model)
    loss_op = F.binary_cross_entropy_with_logits if multi_label else F.cross_entropy
    print(model)
    model.cuda(args.gpu)

    if args.debug_mem:
        print("========== Model, optimizer, and init data===========")
        model.reset_parameters()
        model.train()
        optimizer = get_optimizer(model_config, model)
        optimizer.zero_grad()
        usage = get_memory_usage(args.gpu, True)
        exp_recorder.record("network", 'SAGE')
        exp_recorder.record("model_only", usage / MB, 2)    
        print("========== Load data to GPU ===========")
        for _data in loader:     
            optimizer.zero_grad()
            _data = T.ToSparseTensor()(_data.to('cuda'))
            _data.adj_t.fill_cache_()
            num_nodes, num_edges = _data.num_nodes, _data.num_edges
            print(f'sampled sub-graph: #nodes: {num_nodes}, #edges: {num_edges}')
            init_mem = get_memory_usage(args.gpu, True)
            data_mem = init_mem / MB - exp_recorder.val_dict['model_only']
            exp_recorder.record("data", init_mem / MB - exp_recorder.val_dict['model_only'], 2)
            out = model(_data.x, _data.adj_t)
            y = _data.y
            loss = loss_op(out[_data.train_mask], y[_data.train_mask])
            print("========== Before Backward ===========")
            before_backward = get_memory_usage(args.gpu, True)
            act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss, out])
            res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                            data_mem,
                                                                            act_mem / MB)
            print(res)
            loss.backward()
            optimizer.step()
            del loss
            print("========== After Backward ===========")
            after_backward = get_memory_usage(args.gpu, True)
            total_mem = before_backward + (after_backward - init_mem)
            res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (total_mem / MB,
                                                                            data_mem,
                                                                            act_mem / MB)
            print(res)
            exp_recorder.record("total", total_mem / MB, 2)
            exp_recorder.record("activation", act_mem / MB, 2)
            exit()

    logger = Logger(args.runs, args)
    eval_start_epoch = model_config['eval_start_epoch']
    eval_steps = model_config['eval_steps']
    for run in range(args.runs):
        model.reset_parameters()
        optimizer = get_optimizer(model_config, model)
        for epoch in range(1, 1 + model_config['epochs']):
            loss = train(model, loader, optimizer, loss_op, args.grad_norm)
            if model_config['log_steps'] > 0 and epoch % model_config['log_steps'] == 0:
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Train Loss: {loss:.4f}')
        
            if epoch > eval_start_epoch and epoch % eval_steps == 0:
                result = test(model, data)
                logger.add_result(run, result)
                train_acc, valid_acc, test_acc = result
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Train f1: {100 * train_acc:.2f}%, '
                      f'Valid f1: {100 * valid_acc:.2f}% '
                      f'Test f1: {100 * test_acc:.2f}%')

        logger.add_result(run, result)
        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == '__main__':
    main()
mem_speed_bench/non_ogb_datasets/train_mini_batch.py
saintslab/i-Exact
import argparse
import json
import os
import time


def run_cmd(cmd):
    print(cmd)
    return os.system(cmd)


def to_command(conf_path, batch_size, swap, n_bits, not_compress_act, amp):
    cmd = f"python ./products/train_mini_batch.py --conf {conf_path} --batch_size {batch_size} --debug_mem --gpu 0"
    if not_compress_act:
        cmd += " --not_compress_act"
    if swap:
        cmd += " --swap"
    if amp:
        cmd += " --amp"
    cmd += f" --n_bits {n_bits}"
    return cmd


def round_up(x):
    return int((x + 9)// 10 * 10)

def round_down(x):
    return int(x // 10 * 10)


def run_benchmark(conf_path, batch_size, swap, n_bits, not_compress_act, amp):
    cmd = to_command(conf_path, batch_size, swap, n_bits, not_compress_act, amp)
    ret_code = run_cmd(cmd)

    # if ret_code != 0:
    #     out_file = "speed_results.json"
    #     with open(out_file, "a") as fout:
    #         val_dict = {
    #             "conf_path": conf_path,
    #             "swap": swap,
    #             "batch_size": batch_size,
    #         }
    #         fout.write(json.dumps(val_dict) + "\n")
    #     print(f"save results to {out_file}")

    time.sleep(1)
    run_cmd("nvidia-smi > /dev/null")
    time.sleep(1)
    return ret_code


def binary_search_max_batch(conf_path, swap, n_bits, not_compress_act, amp, low, high):
    ret = 0
    low, high= round_up(low), round_down(high)

    while low <= high:
        mid = round_down(low + (high - low) // 2)
        success = run_benchmark(conf_path, mid, swap, n_bits, not_compress_act, amp) == 0
        if success:
            ret = mid
            low = round_up(mid + 1)
        else:
            high = round_down(mid - 1)

    return ret



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--conf", type=str)
    parser.add_argument("--mode", type=str, default='binary_search_max_batch')
    parser.add_argument("--retry", type=int, default=1)
    args = parser.parse_args()
    if args.mode == 'binary_search_max_batch':
        out_file = "max_batch_results.json"
        low, high = 200000, 1500000
        not_compress_act = True
        max_batch_size = binary_search_max_batch(args.conf, False, 8, True, False, low, high)
        with open(out_file, "a") as fout:
            val_dict = {
                "conf": args.conf,
                "not_compress_act": not_compress_act,
                "max_batch_size": max_batch_size,
                "tstamp": time.time()
            }
            fout.write(json.dumps(val_dict) + "\n")
            print('finished vanilla result')

        for n_bits in [1, 2, 4, 8]:
            for swap in [True, False]:
                for amp in [True, False]:
                    not_compress_act = False
                    max_batch_size = binary_search_max_batch(args.conf, swap, n_bits, not_compress_act, amp, low, high)
                    with open(out_file, "a") as fout:
                        val_dict = {
                            "conf": args.conf,
                            "n_bits": n_bits,
                            "swap": swap,
                            "amp": amp,
                            "max_batch_size": max_batch_size,
                            "tstamp": time.time()
                        }
                        fout.write(json.dumps(val_dict) + "\n")
                    print(f"save results to {out_file}")
    else:
        raise ValueError()

mem_speed_bench/products/exp_mem_speed.py
saintslab/i-Exact
import torch
from tqdm import tqdm
import torch.nn.functional as F

from torch_geometric.nn import SAGEConv, GCNConv, GENConv
from torch.utils.checkpoint import checkpoint

from exact import config


class SharedModule(torch.nn.Module):

    def inference(self, x_all, subgraph_loader):
        pbar = tqdm(total=x_all.size(0) * len(self.convs))
        pbar.set_description('Evaluating')

        for i, conv in enumerate(self.convs):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to('cuda')
                x = x_all[n_id].to('cuda')
                x_target = x[:size[1]]
                x = conv((x, x_target), edge_index)
                if i != len(self.convs) - 1:
                    x = F.relu(x)
                xs.append(x.cpu())

                pbar.update(batch_size)
            del x_all
            x_all = torch.cat(xs, dim=0)
            del xs
        pbar.close()
        return x_all    

    
class GCN(SharedModule):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(GCN, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(
            GCNConv(in_channels, hidden_channels, normalize=False))
        for _ in range(num_layers - 2):
            self.convs.append(
                GCNConv(hidden_channels, hidden_channels, normalize=False))
        self.convs.append(
            GCNConv(hidden_channels, out_channels, normalize=False))
        self.dropout = torch.nn.Dropout(p=dropout)
        self.act = torch.nn.ReLU()

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, adj_t):
        for conv in self.convs[:-1]:
            x = conv(x, adj_t)
            x = self.act(x)
            if not config.debug_remove_dropout:
                x = self.dropout(x)
        x = self.convs[-1](x, adj_t)
        return torch.log_softmax(x, dim=-1)


        
class SAGE(SharedModule):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout, batch_norm=False):
        super(SAGE, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(SAGEConv(in_channels, hidden_channels))
        use_bn = batch_norm and not config.debug_remove_bn
        self.use_bn = use_bn
        if use_bn:
            self.bns = torch.nn.ModuleList()
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
            if use_bn:
                self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.convs.append(SAGEConv(hidden_channels, out_channels))
        self.dropout = torch.nn.Dropout(p=dropout)
        self.act = torch.nn.ReLU()

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        if self.use_bn:
            for bn in self.bns:
                bn.reset_parameters()


    def forward(self, x, edge_index, edge_weight=None):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, edge_index, edge_weight)
            if self.use_bn:
                x = self.bns[i](x)
            x = self.act(x)
            if not config.debug_remove_dropout:
                x = self.dropout(x)
        x = self.convs[-1](x, edge_index, edge_weight)
        return torch.log_softmax(x, dim=-1) 


class DeeperGCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers, mlp_layers, dropout=0.5, 
                 block='res+', gcn_aggr='softmax_sg', norm='batch', t=1., p=1., learn_t=False, learn_p=False,
                 msg_norm=False, learn_msg_scale=False):
        super(DeeperGCN, self).__init__()
        self.checkpoint_grad = False
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.block = block
        self.aggr = gcn_aggr
        self.num_layers = num_layers
        self.mlp_layers = mlp_layers
        self.t, self.p, self.learn_t, self.learn_p = t, p, learn_t, learn_p
        self.msg_norm, self.learn_msg_scale = msg_norm, learn_msg_scale
        self.norm = norm
        self.dropout = torch.nn.Dropout(p=dropout)
        self.act = torch.nn.ReLU()

        if self.aggr in ['softmax_sg', 'softmax', 'power'] and self.num_layers > 3:
            self.checkpoint_grad = True
            self.ckp_k = self.num_layers // 2

        if self.block == 'res+':
            print('BN->ReLU->GraphConv->Res')
        elif self.block == 'res':
            print('GraphConv->BN->ReLU->Res')
        elif self.block == "plain":
            print('GraphConv->LN/BN->ReLU')
        else:
            raise Exception('Unknown block Type')
        if self.norm != 'batch':
            raise NotImplementedError()

        self.gcns = torch.nn.ModuleList()
        self.norms = torch.nn.ModuleList()

        self.node_features_encoder = torch.nn.Linear(in_channels, hidden_channels)
        self.node_pred_linear = torch.nn.Linear(hidden_channels, out_channels)

        for _ in range(self.num_layers):
            gcn = GENConv(hidden_channels, hidden_channels,
                            aggr=self.aggr,
                            t=self.t, learn_t=self.learn_t,
                            p=self.p, learn_p=self.learn_p,
                            msg_norm=self.msg_norm, learn_msg_scale=learn_msg_scale,
                            norm=self.norm, num_layers=self.mlp_layers)

            self.gcns.append(gcn)
            self.norms.append(torch.nn.BatchNorm1d(hidden_channels))
    
    def reset_parameters(self):
        for gcn in self.gcns:
            gcn.reset_parameters()
        for norm in self.norms:
            norm.reset_parameters()
        
    def forward(self,  x, edge_index):
        h = self.node_features_encoder(x)
        if self.block == 'res+':
            h = self.gcns[0](h, edge_index)
            if self.checkpoint_grad:

                for layer in range(1, self.num_layers):
                    h1 = self.norms[layer - 1](h)
                    h2 = self.act(h1)
                    h2 = self.dropout(h2)

                    if layer % self.ckp_k != 0:
                        res = checkpoint(self.gcns[layer], h2, edge_index)
                        h = res + h
                    else:
                        h = self.gcns[layer](h2, edge_index) + h
            else:
                for layer in range(1, self.num_layers):
                    h1 = self.norms[layer - 1](h)
                    h2 = self.act(h1)
                    h2 = self.dropout(h2)
                    h = self.gcns[layer](h2, edge_index) + h

            h = self.act(self.norms[self.num_layers - 1](h))
            h = self.dropout(h)

        elif self.block == 'res':

            h = self.act(self.norms[0](self.gcns[0](h, edge_index)))
            h = self.dropout(h)

            for layer in range(1, self.num_layers):
                h1 = self.gcns[layer](h, edge_index)
                h2 = self.norms[layer](h1)
                h = self.act(h2) + h
                h = self.dropout(h)

        elif self.block == 'plain':

            h = self.act(self.norms[0](self.gcns[0](h, edge_index)))
            h = self.dropout(h)

            for layer in range(1, self.num_layers):
                h1 = self.gcns[layer](h, edge_index)
                h2 = self.norms[layer](h1)
                h = self.act(h2)
                h = self.dropout(h)
        else:
            raise Exception('Unknown block Type')
        h = self.node_pred_linear(h)
        return torch.log_softmax(h, dim=-1)


    def inference(self, x_all, subgraph_loader):
        pbar = tqdm(total=x_all.size(0) * len(self.gcns))
        pbar.set_description('Evaluating')
        if self.block != 'res+':
            raise NotImplementedError
        for i, gcn in enumerate(self.gcns):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to('cuda')
                x = x_all[n_id].to('cuda')
                if i == 0:
                    x = self.node_features_encoder(x)
                    h = gcn(x, edge_index)
                else:
                    h = x
                    h1 = self.norms[i - 1](h)
                    h2 = self.act(h1)
                    h2 = self.dropout(h2)
                    h = gcn(h2, edge_index) + h
                    if i == self.num_layers - 1:
                        h = self.norms[self.num_layers - 1](h)
                        h = F.relu(h)
                        h = self.dropout(h)
                        h = self.node_pred_linear(h)
                xs.append(h.cpu())
                pbar.update(batch_size)
            del x_all
            x_all = torch.cat(xs, dim=0)
            del xs
        pbar.close()
        return x_all    
mem_speed_bench/products/gnn.py
saintslab/i-Exact
import sys
import os

from torch.optim import optimizer
sys.path.append(os.getcwd())
import argparse
import random
import time
import yaml
import math
import warnings

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from torch.cuda.amp import autocast, GradScaler
from torch_sparse import SparseTensor
import torch_geometric.transforms as T
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from ogb.nodeproppred import Evaluator

from exact import get_memory_usage, compute_tensor_bytes, exp_recorder, config, QModule
from utils import cast_int32
import models
from data import get_data
from logger import Logger


EPSILON = 1 - math.log(2)
MB = 1024**2
GB = 1024**3


parser = argparse.ArgumentParser()
parser.add_argument('--conf', type=str, required=True)
parser.add_argument('--root', type=str, default='./data')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--debug_mem', action='store_true',
                    help='whether to debug the memory usage')
parser.add_argument('--n_bits', type=int, default=8)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--simulate', action='store_true')
parser.add_argument('--act_fp', action='store_true')
parser.add_argument('--kept_frac', type=float, default=1.0)
parser.add_argument('--amp', help='whether to enable apx mode', action='store_true')
parser.add_argument('--dropout2', help='whether to use dropout2', action='store_true')


def adjust_learning_rate(optimizer, lr, epoch):
    if epoch <= 50:
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr * epoch / 50


def drop_edge(adj_t, edge_drop):
    if edge_drop == 0:
        return adj_t
    nnz = adj_t.nnz()
    bound = int(nnz * edge_drop)
    perm = torch.randperm(nnz).cuda()
    eids = perm[bound:]
    row, col, value = adj_t.storage._row, adj_t.storage._col, adj_t.storage._value
    row, col= row[eids], col[eids]
    if value is not None:
        value = value[eids]
    adj_t = SparseTensor(row=row, col=col, value=value, 
                            sparse_sizes=(adj_t.size(0), adj_t.size(1)), is_sorted=False)
    return adj_t

def get_optimizer(model_config, model):
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    return optimizer


@torch.no_grad()
def test(model, data, evaluator, amp_mode=False):
    model.eval()
    with autocast(enabled=amp_mode):
        out = model(data.x, data.adj_t)
    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': data.y[data.train_mask],
        'y_pred': y_pred[data.train_mask],
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': data.y[data.val_mask],
        'y_pred': y_pred[data.val_mask],
    })['acc']
    test_acc = evaluator.eval({
        'y_true': data.y[data.test_mask],
        'y_pred': y_pred[data.test_mask],
    })['acc']
    return train_acc, valid_acc, test_acc


def main():
    args = parser.parse_args()
    with open(args.conf, 'r') as fp:
        model_config = yaml.load(fp, Loader=yaml.FullLoader)
    args.model = model_config['name']
    assert args.model.lower() in ['gcn', 'sage', 'gat']
    if args.simulate:
        print('=' * 30 + 'Simulation' + '=' * 30 )
        config.simulate = True
    if args.act_fp and args.kept_frac == 1.0:
        use_qmodule = False
        print('=' * 30 + 'Do not compress activation' + '=' * 30)
        config.compress_activation = False
    else:
        use_qmodule = True
        config.compress_activation = not args.act_fp
        config.activation_compression_bits = [args.n_bits]
        precision = 'FP32' if not config.compress_activation\
                     else f'INT{config.activation_compression_bits[0]}'
        print('=' * 30 + f'Store activation in {precision}' + '=' * 30)
    print(f'use qmodule: {use_qmodule}')
    if args.kept_frac < 1.0:
        config.kept_frac = args.kept_frac
    print(f'kept frac: {args.kept_frac}')
    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        print("Use GPU {} for training".format(args.gpu))
    
    if args.amp:
        config.amp = True
        print(f'amp mode: {config.amp}')
        
    if args.dropout2:
        print('using dropout2')
        config.dropout2 = True

    grad_norm = model_config.get('grad_norm', None)
    print(f'clipping grad norm: {grad_norm}')
    torch.cuda.set_device(args.gpu)

    data, num_features, num_classes = get_data(args.root, 'products')
    evaluator = Evaluator(name='ogbn-products')
    logger = Logger(args.runs, args)
    GNN = getattr(models, args.model)
    model = GNN(in_channels=num_features, out_channels=num_classes, **model_config['architecture'])
    if model_config['loop']:
        t = time.perf_counter()
        print('Adding self-loops...', end=' ', flush=True)
        data.adj_t = data.adj_t.set_diag()
        print(f'Done! [{time.perf_counter() - t:.2f}s]')
    
    if model_config['normalize']:
        t = time.perf_counter()
        print('Normalizing data...', end=' ', flush=True)
        data.adj_t = gcn_norm(data.adj_t, add_self_loops=False)
        print(f'Done! [{time.perf_counter() - t:.2f}s]')

    if use_qmodule:
        print('convert the model')
        model = QModule(model)
    print(model)
    model.cuda(args.gpu)
    if args.debug_mem:
        model.train()
        print("========== Model Only ===========")
        usage = get_memory_usage(args.gpu, True)
        exp_recorder.record("network", 'GCN')
        exp_recorder.record("model_only", usage / MB, 2)
        print("========== Load data to GPU ===========")
        data = T.ToSparseTensor()(data).to('cuda')
        data.adj_t.fill_cache_()
        init_mem = get_memory_usage(args.gpu, True)
        data_mem = init_mem / MB - exp_recorder.val_dict['model_only']
        exp_recorder.record("data", init_mem / MB - exp_recorder.val_dict['model_only'], 2)
        model.reset_parameters()
        model.train()
        optimizer = get_optimizer(model_config, model)
        optimizer.zero_grad()
        out = model(data.x, data.adj_t)[data.train_mask]
        loss = F.nll_loss(out, data.y.squeeze(1)[data.train_mask])
        print("========== Before Backward ===========")
        del out
        before_backward = get_memory_usage(args.gpu, True)
        act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss])
        res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                           data_mem,
                                                                           act_mem / MB)
        print(res) 
        loss.backward()
        optimizer.step()
        del loss
        print("========== After Backward ===========")
        after_backward = get_memory_usage(args.gpu, True)
        total_mem = before_backward + (after_backward - init_mem)
        res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (total_mem / MB,
                                                                           data_mem,
                                                                           act_mem / MB)
        print(res)
        exp_recorder.record("total", total_mem / MB, 2)
        exp_recorder.record("activation", act_mem / MB, 2)
        exp_recorder.dump('mem_results.json')
        exit()
    data = T.ToSparseTensor()(data)
    data = data.to('cuda')
    # ======================================================================
    # Note: if you want to train the model with 32-bit graph structure data,
    # uncomment the below code.

    # cast_int32(data)

    # ======================================================================
    for run in range(args.runs):
        model.reset_parameters()
        optimizer = get_optimizer(model_config, model)
        if args.amp:
            print('activate amp mode')
            scaler = GradScaler()
        else:
            scaler = None
        for epoch in range(1, 1 + model_config['epochs']):
            # ==== train the model ====
            if model_config.get('adjust_lr', False):
                adjust_learning_rate(optimizer, model_config['lr'], epoch)
            model.train()
            optimizer.zero_grad()
            with autocast(enabled=args.amp):
                out = model(data.x, data.adj_t)
                loss = F.cross_entropy(out[data.train_mask], data.y.squeeze(1)[data.train_mask])
            torch.cuda.empty_cache()
            if args.amp:
                scaler.scale(loss).backward()
                if grad_norm is not None:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                if grad_norm is not None:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
                optimizer.step()
            # ===========================
            result = test(model, data, evaluator, args.amp)
            logger.add_result(run, result)
            if  model_config['log_steps'] > 0 and epoch % model_config['log_steps'] == 0:
                train_acc, valid_acc, test_acc = result
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Loss: {loss.item():.4f}, '
                      f'Train: {100 * train_acc:.2f}%, '
                      f'Valid: {100 * valid_acc:.2f}% '
                      f'Test: {100 * test_acc:.2f}%')

        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == '__main__':
    main()
mem_speed_bench/products/train_full_batch.py
saintslab/i-Exact
import sys
import os
sys.path.append(os.getcwd())
import argparse
import random
import time
import warnings
from torch_geometric.data.cluster import ClusterLoader
import yaml
import pdb

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from torch.cuda.amp import autocast, GradScaler

from torch_geometric.utils import subgraph
from torch_geometric.data import GraphSAINTRandomWalkSampler, NeighborSampler, ClusterData

from exact import get_memory_usage, compute_tensor_bytes, exp_recorder, config, QModule, cast_adj
import models
from data import get_data
from logger import Logger

import torch_geometric.transforms as T
from ogb.nodeproppred import Evaluator

MB = 1024**2
GB = 1024**3

SAMPLER = {'graph_saint': GraphSAINTRandomWalkSampler, 'cluster_gcn': ClusterData}

parser = argparse.ArgumentParser()
parser.add_argument('--conf', type=str, required=True, 
                    help='the path to the configuration file')
parser.add_argument('--batch_size', type=int, required=True)
parser.add_argument('--root', type=str, default='./data')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=0, type=int,
                    help='GPU id to use.')
parser.add_argument('--num_workers', type=int, default=12)
parser.add_argument('--debug_mem', action='store_true',
                    help='whether to debug the memory usage')
parser.add_argument('--n_bits', type=int, default=8)
parser.add_argument('--runs', type=int, default=10)
parser.add_argument('--simulate', action='store_true')
parser.add_argument('--act_fp', action='store_true')
parser.add_argument('--kept_frac', type=float, default=1.0)
parser.add_argument('--amp', help='whether to enable apx mode', action='store_true')
parser.add_argument('--dropout2', help='whether to use dropout2', action='store_true')



def get_optimizer(model_config, model):
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    if model_config['optim'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=model_config['lr'])
    elif model_config['optim'] == 'rmsprop':
        optimizer = torch.optim.RMSprop(model.parameters(), lr=model_config['lr'])
    else:
        raise NotImplementedError
    return optimizer


def train(model, loader, optimizer, scaler, amp_mode, grad_norm):
    model.train()

    total_loss = 0
    for data in loader:
        data = T.ToSparseTensor()(data.to('cuda'))
        optimizer.zero_grad()
        with autocast(enabled=amp_mode):
            out = model(data.x, data.adj_t)
            y = data.y.squeeze(1)
            loss = F.cross_entropy(out[data.train_mask], y[data.train_mask])
        del data
        if scaler is not None:
            scaler.scale(loss).backward()
            if grad_norm is not None:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            if grad_norm is not None:
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
            optimizer.step()
        total_loss += loss.item()

    return total_loss / len(loader)


@torch.no_grad()
def test(model, data, evaluator, subgraph_loader, amp_mode):
    model.eval()
    with autocast(enabled=amp_mode):
        out = model.mini_inference(data.x, subgraph_loader)

    y_true = data.y
    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': y_true[data.train_mask],
        'y_pred': y_pred[data.train_mask]
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': y_true[data.val_mask],
        'y_pred': y_pred[data.val_mask]
    })['acc']
    test_acc = evaluator.eval({
        'y_true': y_true[data.test_mask],
        'y_pred': y_pred[data.test_mask]
    })['acc']

    return train_acc, valid_acc, test_acc


def main():
    args = parser.parse_args()
    with open(args.conf, 'r') as fp:
        model_config = yaml.load(fp, Loader=yaml.FullLoader)
    args.model = model_config['arch_name']
    assert model_config['name'] in ['graph_saint', 'cluster_gcn']
    if args.simulate:
        print('=' * 30 + 'Simulation' + '=' * 30 )
        config.simulate = True
    if args.act_fp and args.kept_frac == 1.0:
        use_qmodule = False
        print('=' * 30 + 'Do not compress activation' + '=' * 30)
        config.compress_activation = False
    else:
        use_qmodule = True
        config.compress_activation = not args.act_fp
        config.activation_compression_bits = [args.n_bits]
        precision = 'FP32' if not config.compress_activation\
                     else f'INT{config.activation_compression_bits[0]}'
        print('=' * 30 + f'Store activation in {precision}' + '=' * 30)
    print(f'use qmodule: {use_qmodule}')
    if args.kept_frac < 1.0:
        config.kept_frac = args.kept_frac
    print(f'kept frac: {args.kept_frac}')
    print(f'amp mode: {config.amp}')
    if args.amp:
        print('activate amp mode...')
        config.amp = True
        scaler = GradScaler()
    else:
        config.amp = False
        scaler = None

    if args.dropout2:
        print('using dropout2')
        config.dropout2 = True

    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        print("Use GPU {} for training".format(args.gpu))
    grad_norm = model_config.get('grad_norm', None)
    print(f'clipping grad norm: {grad_norm}')
    torch.cuda.set_device(args.gpu)
    data, num_features, num_classes = get_data(args.root, 'products')
    sampler_cls = SAMPLER[model_config['name']]
    s_time = time.time()
    print("=" * 30 + f'Prepare {sampler_cls.__name__} for training' + '=' * 30)
    if model_config['name'] == 'graph_saint':
        loader = sampler_cls(data,
                            batch_size=args.batch_size,
                            num_workers=args.num_workers,
                            **model_config['sampler'])

    elif model_config['name'] == 'cluster_gcn':
        cluster_data = sampler_cls(data, 
                                   **model_config['sampler'])
        loader = ClusterLoader(cluster_data, batch_size=args.batch_size, 
                               shuffle=True, num_workers=args.num_workers)
    else:
        raise NotImplementedError
    print("=" * 30 + 
          f'Finished Building {sampler_cls.__name__}, used {time.time() - s_time} sec' + 
          '=' * 30)
    GNN = getattr(models, args.model)
    model = GNN(in_channels=num_features, out_channels=num_classes, **model_config['architecture'])
    if use_qmodule:
        print('convert the model')
        model = QModule(model)
    print(model)
    model.cuda(args.gpu)
    evaluator = Evaluator(name='ogbn-products')
    logger = Logger(args.runs, args)
    if args.debug_mem:
        print("========== Model Only ===========")
        usage = get_memory_usage(args.gpu, True)
        exp_recorder.record("network", 'SAGE')
        exp_recorder.record("model_only", usage / MB, 2)    
        model.reset_parameters()
        model.train()
        optimizer = get_optimizer(model_config, model)
        print("========== Load data to GPU ===========")
        for data in loader:     
            optimizer.zero_grad()
            data = T.ToSparseTensor()(data.to('cuda'))
            data.adj_t.fill_cache_()
            num_nodes, num_edges = data.num_nodes, data.num_edges
            print(f'sampled sub-graph: #nodes: {num_nodes}, #edges: {num_edges}')
            init_mem = get_memory_usage(args.gpu, True)
            data_mem = init_mem / MB - exp_recorder.val_dict['model_only']
            exp_recorder.record("data", init_mem / MB - exp_recorder.val_dict['model_only'], 2)
            with autocast(enabled=args.amp):
                out = model(data.x, data.adj_t)
                y = data.y.squeeze(1)
                loss = F.cross_entropy(out[data.train_mask], y[data.train_mask])
            print("========== Before Backward ===========")
            before_backward = get_memory_usage(args.gpu, True)
            act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss, out])
            res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                            data_mem,
                                                                            act_mem / MB)
            print(res)
            if args.amp:
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                optimizer.step()

            del loss
            print("========== After Backward ===========")
            after_backward = get_memory_usage(args.gpu, True)
            total_mem = before_backward + (after_backward - init_mem)
            res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (total_mem / MB,
                                                                            data_mem,
                                                                            act_mem / MB)
            print(res)
            exp_recorder.record("total", total_mem / MB, 2)
            exp_recorder.record("activation", act_mem / MB, 2)
            exp_recorder.dump('mem_results.json')
            exit()

    print("=" * 30 + 'Prepare NeighborSampler for evaluation' + '=' * 30)
    s_time = time.time()
    subgraph_loader = NeighborSampler(data.edge_index, sizes=[-1],
                                      batch_size=4096, shuffle=False)
    print("=" * 30 + 
          f'Finished Building NeighborSampler, used {time.time() - s_time} sec' + 
          '=' * 30)
    data.edge_index = None
    eval_start_epoch = model_config['eval_start_epoch']
    eval_steps = model_config['eval_steps']

    for run in range(args.runs):
        model.reset_parameters()
        optimizer = get_optimizer(model_config, model)
        for epoch in range(1, 1 + model_config['epochs']):
            loss = train(model, loader, optimizer, scaler, amp_mode=config.amp, grad_norm=grad_norm)
            if model_config['log_steps'] > 0 and epoch % model_config['log_steps'] == 0:
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Train Loss: {loss:.4f}')

            if epoch > eval_start_epoch and epoch % eval_steps == 0:
                result = test(model, data, evaluator, subgraph_loader, amp_mode=args.amp)
                logger.add_result(run, result)
                train_acc, valid_acc, test_acc = result
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Train: {100 * train_acc:.2f}%, '
                      f'Valid: {100 * valid_acc:.2f}% '
                      f'Test: {100 * test_acc:.2f}%')

        logger.add_result(run, result)
        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == '__main__':
    main()
mem_speed_bench/products/train_mini_batch.py
saintslab/i-Exact

from typing import Optional, Tuple

import torch
from torch import Tensor
import torch.nn.functional as F
from exact import get_memory_usage
from torch_sparse import SparseTensor
from ogb.nodeproppred import Evaluator

def index2mask(idx: Tensor, size: int) -> Tensor:
    mask = torch.zeros(size, dtype=torch.bool, device=idx.device)
    mask[idx] = True
    return mask

# from torch_sparse.storage import SparseStorage
def cast_int32(data):
    data.adj_t.storage.rowptr()
    data.adj_t.storage.col()
    data.adj_t.storage.row()
    data.adj_t.storage.rowcount() 
    data.adj_t.storage.colptr()
    data.adj_t.storage._colptr = data.adj_t.storage._colptr.int()
    data.adj_t.storage._rowptr = data.adj_t.storage._rowptr.int()
    data.adj_t.storage._col = data.adj_t.storage._col.int()
    data.adj_t.storage._row = data.adj_t.storage._row.int()
    data.adj_t.storage._rowcount = data.adj_t.storage._rowcount.int()
mem_speed_bench/utils.py
saintslab/i-Exact
import pdb

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch_geometric.nn.norm import msg_norm

import torch_sparse
from torch_geometric.nn import GCNConv, SAGEConv, GCN2Conv, GENConv


from exact import config, QGCNConv, QSAGEConv, QGCN2Conv, QGENConv
from exact import get_memory_usage, compute_tensor_bytes


config.activation_compression_bits = [2]
num_samples = 10000
in_channels, out_channels = 100, 128
nnz = 50000
i = torch.randint(high=num_samples, size=(2, nnz), dtype=torch.long)
adj_t = torch_sparse.SparseTensor.from_edge_index(i, sparse_sizes=[num_samples, num_samples]).cuda()
adj_t.fill_cache_()


def test_gcnconv_correctness():
    layer = GCNConv(in_channels, out_channels).cuda()
    qlayer = QGCNConv(in_channels, out_channels).cuda()
    x = torch.randn((num_samples, in_channels)).cuda()
    y = torch.empty(num_samples, dtype=torch.long).random_(out_channels).cuda()
    print('==================test GCNConv correctness==================')
    with torch.no_grad():
        qlayer.weight.copy_(layer.weight.t())

    ce = nn.CrossEntropyLoss().cuda()

    def get_grad(model):
        pred = model(x, adj_t)
        pred = F.relu(pred)
        loss = ce(pred, y)
        model.weight.grad = None
        model.bias.grad = None
        loss.backward()
        return model.weight.grad.cpu().numpy()

    true_grad = get_grad(layer).T
    grads = []
    for i in range(10):
        grads.append(get_grad(qlayer))
    grads = np.stack(grads, 0)
    grad_mean = grads.mean(0)
    grad_std = grads.std(0)
    bias = np.linalg.norm(grad_mean - true_grad)
    print('Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(true_grad), bias, np.linalg.norm(grad_std)))


def test_gcnconv_memory():
    print("========== Gcnconv Memory Test ==========")
    adj_t_ = adj_t.set_diag()
    deg = adj_t_.sum(dim=1).to(torch.float)
    deg_inv_sqrt = deg.pow(-0.5)
    deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
    adj_t_ = deg_inv_sqrt.view(-1, 1) * adj_t_ * deg_inv_sqrt.view(1, -1)

    layer = GCNConv(in_channels, in_channels, cached=False, normalize=False).cuda()
    qlayer = QGCNConv(in_channels, in_channels, cached=False, normalize=False).cuda()
    x = torch.randn((num_samples, in_channels))
    with torch.no_grad():
        qlayer.weight.copy_(layer.weight.t())


    def test_implementation(model, n_layers):
        data = x.to("cuda").requires_grad_()
        output = data
        before = get_memory_usage(0)
        for _ in range(n_layers):
            output = model(output, adj_t_)
        after = get_memory_usage(0) - compute_tensor_bytes([output])
        if model.__class__.__name__ == 'GCNConv':
            after += compute_tensor_bytes([data])
        return after - before

    usage_ref = test_implementation(layer, 5)
    usage_us = test_implementation(qlayer, 5)
    print("5 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("5 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("5 layer: Ratio: %.2f" % (usage_ref / usage_us))

    usage_ref = test_implementation(layer, 10)
    usage_us = test_implementation(qlayer, 10)
    print("10 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("10 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("10 layer: Ratio: %.2f" % (usage_ref / usage_us))



def test_sageconv_correctness():
    layer = SAGEConv(in_channels, out_channels, root_weight=True, bias=True).cuda()
    qlayer = QSAGEConv(in_channels, out_channels, root_weight=True, bias=True).cuda()
    x = torch.randn((num_samples, in_channels)).cuda()
    y = torch.empty(num_samples, dtype=torch.long).random_(out_channels).cuda()
    print('==================test SAGEConv correctness==================')
    with torch.no_grad():
        qlayer.lin_l.weight.copy_(layer.lin_l.weight)
        qlayer.lin_r.weight.copy_(layer.lin_r.weight)
        qlayer.lin_l.bias.copy_(layer.lin_l.bias)

    ce = nn.CrossEntropyLoss().cuda()

    def get_grad(model):
        pred = model(x, adj_t)
        pred = F.relu(pred)
        loss = ce(pred, y)
        model.lin_l.weight.grad = None
        model.lin_r.weight.grad = None
        loss.backward()
        return [lin.weight.grad.cpu().numpy() for lin in [model.lin_l, model.lin_r]]

    lin_l_true_grad, lin_r_true_grad = get_grad(layer)
    lin_l_grads, lin_r_grads = [], []
    for i in range(10):
        lin_l_grads.append(get_grad(qlayer)[0])
        lin_r_grads.append(get_grad(qlayer)[1])
    
    for idx, (true_grad, grads) in enumerate(zip([lin_l_true_grad, lin_r_true_grad], [lin_l_grads, lin_r_grads])):
        if idx == 0:
            prefix = 'lin l '
        else:
            prefix = 'lin r '
        grads = np.stack(grads, 0)
        grad_mean = grads.mean(0)
        grad_std = grads.std(0)
        bias = np.linalg.norm(grad_mean - true_grad)
        print(prefix + 'Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(true_grad), bias, np.linalg.norm(grad_std)))


def test_sageconv_memory():
    print("========== SAGEconv Memory Test ==========")
    root_weight = True
    layer = SAGEConv(in_channels, in_channels, root_weight=root_weight, bias=True).cuda()
    qlayer = QSAGEConv(in_channels, in_channels, root_weight=root_weight, bias=True).cuda()
    x = torch.randn((num_samples, in_channels))
    with torch.no_grad():
        qlayer.lin_l.weight.copy_(layer.lin_l.weight)
        qlayer.lin_l.bias.copy_(layer.lin_l.bias)
        if root_weight:
            qlayer.lin_r.weight.copy_(layer.lin_r.weight)


    def test_implementation(model, n_layers):
        data = x.to("cuda").requires_grad_()
        output = data
        before = get_memory_usage(0)
        for _ in range(n_layers):
            output = model(output, adj_t)
        after = get_memory_usage(0) - compute_tensor_bytes([output])
        if model.__class__.__name__ == 'SAGEConv':
            after += compute_tensor_bytes([data])
        return after - before

    usage_ref = test_implementation(layer, 5)
    usage_us = test_implementation(qlayer, 5)
    print("5 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("5 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("5 layer: Ratio: %.2f" % (usage_ref / usage_us))

    usage_ref = test_implementation(layer, 10)
    usage_us = test_implementation(qlayer, 10)
    print("10 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("10 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("10 layer: Ratio: %.2f" % (usage_ref / usage_us))


def test_genconv_correctness():
    layer = GENConv(in_channels, out_channels, num_layers=1).cuda()
    qlayer = QGENConv(in_channels, out_channels, num_layers=1).cuda()
    # pdb.set_trace()
    x = torch.randn((num_samples, in_channels)).cuda()
    y = torch.empty(num_samples, dtype=torch.long).random_(out_channels).cuda()
    print('==================test SAGEConv correctness==================')
    with torch.no_grad():
        for lin, qlin in zip(layer.mlp, qlayer.mlp):
            qlin.weight.copy_(lin.weight)
            qlin.bias.copy_(lin.bias)


    ce = nn.CrossEntropyLoss().cuda()

    def get_grad(model):
        pred = model(x, adj_t)
        pred = F.relu(pred)
        loss = ce(pred, y)
        for lin in model.mlp:
            lin.weight.grad = None
            lin.bias.grad = None
        loss.backward()
        return model.mlp[0].weight.grad.cpu().numpy()

    true_grad = get_grad(layer)
    grads = []
    for i in range(10):
        grads.append(get_grad(qlayer))
    grads = np.stack(grads, 0)
    grad_mean = grads.mean(0)
    grad_std = grads.std(0)
    bias = np.linalg.norm(grad_mean - true_grad)
    print('Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(true_grad), bias, np.linalg.norm(grad_std)))


def test_genconv_memory():
    print("========== GENConv Memory Test ==========")
    layer = GENConv(in_channels*5, in_channels*5, num_layers=3, aggr='softmax_sg').cuda()
    qlayer = QGENConv(in_channels*5, in_channels*5, num_layers=3, aggr='softmax_sg').cuda()
    x = torch.randn((num_samples, in_channels*5))
    with torch.no_grad():
        for lin, qlin in zip(layer.mlp, qlayer.mlp):
            if not isinstance(lin, torch.nn.Linear):
                continue
            qlin.weight.copy_(lin.weight)
            qlin.bias.copy_(lin.bias)
    def test_implementation(model, n_layers):
        data = x.to("cuda").requires_grad_()
        output = data
        before = get_memory_usage(0)
        for _ in range(n_layers):
            output = model(output, adj_t)
        after = get_memory_usage(0) - compute_tensor_bytes([output])
        if model.__class__.__name__ == 'SAGEConv':
            after += compute_tensor_bytes([data])
        return after - before


    usage_ref = test_implementation(layer, 5)
    usage_us = test_implementation(qlayer, 5)
    print("5 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("5 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("5 layer: Ratio: %.2f" % (usage_ref / usage_us))

    usage_ref = test_implementation(layer, 10)
    usage_us = test_implementation(qlayer, 10)
    print("10 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("10 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("10 layer: Ratio: %.2f" % (usage_ref / usage_us))

if __name__ == '__main__':
    test_gcnconv_correctness()
    test_gcnconv_memory()
    test_sageconv_correctness()
    test_sageconv_memory()
    test_genconv_correctness()
    test_genconv_memory()
test/test_conv.py
saintslab/i-Exact
"""Test the activation quantized ops"""

import copy
import numpy as np
from numpy import random
import torch
import torch_sparse
from torch.nn import functional as F


from timeit_v2 import py_benchmark

import exact.cpp_extension.quantization as ext_quantization
from exact import get_memory_usage, compute_tensor_bytes
from exact.conf import config
from exact.ops import qlinear, qbatch_norm, qelu
from exact.spmm import qmatmul

config.activation_compression_bits = [2]
config.compress_activation = True

def test_relu_correctness():
    print("========== ReLU Correctness Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(10000, 128).astype(dtype)

        def test_implementation(func):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            output = func(data)
            output.backward(torch.ones_like(output))

            return [x.detach().cpu().numpy() for x in [output, data.grad]]

        output_ref, grad_data_ref = test_implementation(F.relu)
        output_us, grad_data_us = test_implementation(ext_quantization.act_quantized_relu)

        np.testing.assert_allclose(output_ref, output_us)
        np.testing.assert_allclose(grad_data_ref, grad_data_us)


def test_relu_memory():
    print("========== ReLU Memory Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(10000, 128).astype(dtype)

        def test_implementation(func):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            before = get_memory_usage(0)

            for i in range(10):
                data = func(data)

            after = get_memory_usage(0)

            return after - before

        usage_ref = test_implementation(F.relu)
        usage_us = test_implementation(ext_quantization.act_quantized_relu)

        print("Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
        print("Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))


def test_relu_speed():
    print("========== ReLU Speed Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")

        data_np = np.random.randn(100000, 128).astype(dtype)

        def test_implementation(func):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            stmt = "func(data)"
            t_forward = py_benchmark(stmt, {**globals(), **locals()},
                                     setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            output = func(data)
            head = torch.ones_like(output)
            stmt = "output.backward(head, retain_graph=True)"
            t_backward = py_benchmark(stmt, {**globals(), **locals()},
                                      setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            return t_forward, t_backward

        forward_ref, backward_ref = test_implementation(F.relu)
        forward_us, backward_us = test_implementation(ext_quantization.act_quantized_relu)

        print("Exact.     forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_ref * 1e3, backward_ref * 1e3, (forward_ref + backward_ref) * 1e3))
        print("Quantized. forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_us * 1e3, backward_us * 1e3, (forward_us + backward_us) * 1e3))

def test_dropout_speed():
    print("========== Dropout Speed Test ==========")
    p = 0.5
    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(100000, 128).astype(dtype)

        def test_implementation(func, p):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            stmt = "func(data, p, True)"
            t_forward = py_benchmark(stmt, {**globals(), **locals()},
                                     setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            output = func(data, p, True)
            head = torch.ones_like(output)
            stmt = "output.backward(head, retain_graph=True)"
            t_backward = py_benchmark(stmt, {**globals(), **locals()},
                                      setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            return t_forward, t_backward

        forward_ref, backward_ref = test_implementation(F.dropout, p)
        forward_us, backward_us = test_implementation(ext_quantization.act_quantized_dropout, p)

        print("Exact.     forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_ref * 1e3, backward_ref * 1e3, (forward_ref + backward_ref) * 1e3))
        print("Quantized. forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_us * 1e3, backward_us * 1e3, (forward_us + backward_us) * 1e3))

def test_linear_speed():
    print("========== Linear Speed Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")

        data_np = np.random.randn(100000, 128).astype(dtype)
        w_np = np.random.randn(128, 128).astype(dtype)
        b_np = np.random.randn(128).astype(dtype)
        def test_implementation(func):
            data = torch.tensor(data_np).to("cuda").requires_grad_()
            weight = torch.tensor(w_np).to("cuda").requires_grad_()
            bias = torch.tensor(b_np).to("cuda").requires_grad_()
            stmt = "func(data, weight, bias)"
            t_forward = py_benchmark(stmt, {**globals(), **locals()},
                                     setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            output = func(data, weight, bias)
            head = torch.ones_like(output)
            stmt = "output.backward(head, retain_graph=True)"
            t_backward = py_benchmark(stmt, {**globals(), **locals()},
                                      setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            return t_forward, t_backward

        forward_ref, backward_ref = test_implementation(F.linear)
        forward_us, backward_us = test_implementation(qlinear.apply)

        print("Exact.     forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_ref * 1e3, backward_ref * 1e3, (forward_ref + backward_ref) * 1e3))
        print("Quantized. forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_us * 1e3, backward_us * 1e3, (forward_us + backward_us) * 1e3))


def test_leaky_relu_correctness():
    print("========== LeakyReLU Correctness Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(10000, 128).astype(dtype)

        def test_implementation(func, neg_slope):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            output = func(data, neg_slope)
            output.backward(torch.ones_like(output))

            return [x.detach().cpu().numpy() for x in [output, data.grad]]

        output_ref, grad_data_ref = test_implementation(F.leaky_relu, 0.2)
        output_us, grad_data_us = test_implementation(ext_quantization.act_quantized_leaky_relu, 0.2)

        np.testing.assert_allclose(output_ref, output_us)
        np.testing.assert_allclose(grad_data_ref, grad_data_us)


def test_leaky_relu_memory():
    print("========== LeakyReLU Memory Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(100000, 128).astype(dtype)

        def test_implementation(func, neg_slope):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            before = get_memory_usage(0)

            for i in range(10):
                data = func(data, neg_slope)

            after = get_memory_usage(0)

            return after - before

        usage_ref = test_implementation(F.leaky_relu, 0.2)
        usage_us = test_implementation(ext_quantization.act_quantized_leaky_relu, 0.2)

        print("Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
        print("Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))


def test_leaky_relu_speed():
    print("========== LeakyReLU Speed Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")

        data_np = np.random.randn(10000, 128).astype(dtype)

        def test_implementation(func, neg_slope):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            stmt = "func(data, neg_slope)"
            t_forward = py_benchmark(stmt, {**globals(), **locals()},
                                     setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            output = func(data, neg_slope)
            head = torch.ones_like(output)
            stmt = "output.backward(head, retain_graph=True)"
            t_backward = py_benchmark(stmt, {**globals(), **locals()},
                                      setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

            return t_forward, t_backward

        forward_ref, backward_ref = test_implementation(F.leaky_relu, 0.2)
        forward_us, backward_us = test_implementation(ext_quantization.act_quantized_leaky_relu, 0.2)

        print("Exact.     forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_ref * 1e3, backward_ref * 1e3, (forward_ref + backward_ref) * 1e3))
        print("Quantized. forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
              (forward_us * 1e3, backward_us * 1e3, (forward_us + backward_us) * 1e3))


def test_elu_correctness():
    print("========== ELU Correctness Test ==========")

    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(100, 128).astype(dtype)
        alpha = 0.2
        def test_implementation(func):
            data = torch.tensor(data_np).to("cuda").requires_grad_()

            output = func(data, alpha)
            output.backward(torch.ones_like(output))
            return [x.detach().cpu().numpy() for x in [output, data.grad]]

        output_ref, grad_data_ref = test_implementation(F.elu)
        outputs, grads = [], []
        for _ in range(10):
            output_us, grad_data_us = test_implementation(qelu.apply)
            np.testing.assert_allclose(output_ref, output_us)
            grads.append(grad_data_us)
        grads = np.stack(grads, 0)
        grad_mean = grads.mean(0)
        grad_std = grads.std(0)
        bias = np.linalg.norm(grad_mean - grad_data_ref)
        print('Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(grad_data_ref), bias, np.linalg.norm(grad_std)))


def test_elu_memory():
    print("========== ELU Memory Test ==========")
    alpha = 0.0
    for dtype in ['float32']:
        print(f"test {dtype}...")
        data_np = np.random.randn(10000, 128).astype(dtype)
        def test_implementation(func):
            data = torch.tensor(data_np).to("cuda").requires_grad_()
            output = data
            before = get_memory_usage(0)

            for i in range(10):
                output = func(output, alpha)
            after = get_memory_usage(0) - compute_tensor_bytes(output)
            if func == F.elu:
                after += compute_tensor_bytes(data)
            return after - before

        usage_ref = test_implementation(F.elu)
        usage_us = test_implementation(qelu.apply)
        print("Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
        print("Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
        print("Ratio: %.2f" % (usage_ref / usage_us))


def test_linear_correctness():
    print("========== Linear Correctness Test ==========")
    data_np = np.random.randn(100000, 128).astype('float32')
    data = torch.tensor(data_np).to("cuda")
    # data = arxiv_data.x.continguous()
    ce = torch.nn.CrossEntropyLoss().cuda()
    y = torch.empty(100000, dtype=torch.long).random_(4).cuda()

    def test_implementation(func, weight, bias):
        pred = func(data, weight, bias)
        pred = F.relu(pred)
        pred = pred.reshape(pred.shape[0], 4, pred.shape[1] // 4).mean(2)
        loss = ce(pred, y)
        weight.grad = None
        bias.grad = None
        loss.backward()
        return weight.grad.cpu().numpy()

    w = torch.randn((128, 128), requires_grad=True, device='cuda')
    b = torch.randn((128,), requires_grad=True, device='cuda')
    qw = torch.randn((128, 128), requires_grad=True, device='cuda')
    qb = torch.randn((128,), requires_grad=True, device='cuda')
    with torch.no_grad():
        qw.copy_(w)
        qb.copy_(b)
    true_grad = test_implementation(F.linear, w, b)
    grads = []
    for i in range(10):
        grads.append(test_implementation(qlinear.apply, qw, qb))
    grads = np.stack(grads, 0)
    grad_mean = grads.mean(0)
    grad_std = grads.std(0)
    bias = np.linalg.norm(grad_mean - true_grad)
    print('Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(true_grad), bias, np.linalg.norm(grad_std)))


def test_linear_memory():
    print("========== Linear Memory Test ==========")
    data_np = np.random.randn(169343, 128).astype('float32')
    w = torch.randn((128, 128), requires_grad=True, device='cuda')
    b = torch.randn((128,), requires_grad=True, device='cuda')
    qw = torch.randn((128, 128), requires_grad=True, device='cuda')
    qb = torch.randn((128,), requires_grad=True, device='cuda')
    with torch.no_grad():
        qw.copy_(w)
        qb.copy_(b)

    def test_implementation(func, weight, bias, n_layers):
        data = torch.tensor(data_np).to("cuda").requires_grad_()
        output = data

        before = get_memory_usage(0)

        for i in range(n_layers):
            output = func(output, weight, bias)

        after = get_memory_usage(0) - compute_tensor_bytes([output])
        if func == F.linear:
            after += compute_tensor_bytes([data])

        return after - before

    usage_ref = test_implementation(F.linear, w, b, 5)
    usage_us = test_implementation(qlinear.apply, qw, qb, 5)
    print("5 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("5 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("5 layer: Ratio: %.2f" % (usage_ref / usage_us))
    print("")

    usage_ref = test_implementation(F.linear, w, b, 10)
    usage_us = test_implementation(qlinear.apply, qw, qb, 10)
    print("10 layer: Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("10 layer: Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
    print("10 layer: Ratio: %.2f" % (usage_ref / usage_us))
    print("")


def test_bn_correctness():
    # arguments and test data
    N, CI = 160000, 128
    data_np = np.random.randn(N, CI).astype('float32') * 0.01
    running_mean_np = np.random.randn(CI).astype('float32')
    running_var_np = np.random.randn(CI).astype('float32')
    bn_weight_np = np.random.randn(CI).astype('float32')
    bn_bias_np = np.random.randn(CI).astype('float32')
    training = True

    bn_scheme = None
    config.compress_activation = False

    def test_implementation(func):
        torch.manual_seed(0)
        data = torch.tensor(data_np).to("cuda").requires_grad_()
        running_mean = torch.tensor(running_mean_np).to("cuda")
        running_var = torch.tensor(running_var_np).to("cuda")
        bn_weight = torch.tensor(bn_weight_np).to("cuda").requires_grad_()
        bn_bias = torch.tensor(bn_bias_np).to("cuda").requires_grad_()

        if func == F.batch_norm:
            output = func(data, running_mean, running_var, bn_weight, bn_bias, training, 0.1, 1e-5)
        else:
            output = func(data, running_mean, running_var, bn_weight, bn_bias, training, 0.1, 1e-5, bn_scheme)

        output.backward(torch.ones_like(output))

        return [x.detach().cpu().numpy() for x in [output, data.grad, bn_weight.grad, bn_bias.grad]]

    output_ref, grad_data_ref, grad_weight_ref, grad_bias_ref = test_implementation(F.batch_norm)
    output_us, grad_data_us, grad_weight_us, grad_bias_us = test_implementation(qbatch_norm.apply)

    atol = 1e-3
    rtol = 1e-3
    print("========== BN Correctness Test ==========")
    np.testing.assert_allclose(output_ref, output_us, atol=atol, rtol=rtol)
    np.testing.assert_allclose(grad_data_ref, grad_data_us, atol=atol, rtol=rtol)
    np.testing.assert_allclose(grad_weight_ref, grad_weight_us, atol=atol, rtol=rtol)
    np.testing.assert_allclose(grad_bias_ref, grad_bias_us, atol=atol, rtol=rtol)


def test_bn_speed():
    print("========== Batch norm Speed Test ==========")
    N, CI = 100000, 128
    data_np = np.random.randn(N, CI).astype('float32') * 0.01
    running_mean_np = np.random.randn(CI).astype('float32')
    running_var_np = np.random.randn(CI).astype('float32')
    bn_weight_np = np.random.randn(CI).astype('float32')
    bn_bias_np = np.random.randn(CI).astype('float32')
    config.compress_activation = False

    def test_implementation(func):
        # torch.manual_seed(0)
        data = torch.tensor(data_np).to("cuda").requires_grad_()
        running_mean = torch.tensor(running_mean_np).to("cuda")
        running_var = torch.tensor(running_var_np).to("cuda")
        bn_weight = torch.tensor(bn_weight_np).to("cuda").requires_grad_()
        bn_bias = torch.tensor(bn_bias_np).to("cuda").requires_grad_()
        training = True
        a, b = 0.1, 1e-5
        bn_scheme = None
        if func == F.batch_norm:
            stmt = "func(data, running_mean, running_var, bn_weight, bn_bias, training, a, b)"
        else:
            stmt = "func(data, running_mean, running_var, bn_weight, bn_bias, training, a, b, bn_scheme)"
        t_forward = py_benchmark(stmt, {**globals(), **locals()},
                                 setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

        if func == F.batch_norm:
            output = func(data, running_mean, running_var, bn_weight, bn_bias, training, a, b)
        else:
            output = func(data, running_mean, running_var, bn_weight, bn_bias, training, a, b, None)
        tmp = torch.ones_like(output)
        stmt = "output.backward(tmp, retain_graph=True)"
        t_backward = py_benchmark(stmt, {**globals(), **locals()},
                                  setup="torch.cuda.synchronize()", finish="torch.cuda.synchronize()")

        return t_forward, t_backward

    forward_ref, backward_ref = test_implementation(F.batch_norm)
    forward_us, backward_us = test_implementation(qbatch_norm.apply)

    print("Exact.     forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
          (forward_ref * 1e3, backward_ref * 1e3, (forward_ref + backward_ref) * 1e3))
    print("Quantized. forward: %.2f ms\tbackward: %.2f ms\tsum: %.2f ms" %
          (forward_us * 1e3, backward_us * 1e3, (forward_us + backward_us) * 1e3))


def test_dropout_correctness():
    print("========== Dropout Correctness Test ==========")
    data_np = np.ones((232965, 602)).astype('float32')
    data = torch.tensor(data_np).to("cuda").requires_grad_()
    p = 0.5
    for training in [True]:
        output_ref = F.dropout(data, p, training)
        output_ours = ext_quantization.act_quantized_dropout(data, p, training)
        output_ref.backward(torch.ones_like(output_ref))
        grad_ref = copy.deepcopy(data.grad.detach().cpu().numpy())
        data.grad = None
        output_ours.backward(torch.ones_like(output_ours))
        grad_ours = data.grad.detach().cpu().numpy()
        tmp1 = (output_ours == 0.0).type(torch.float).cpu().numpy()
        tmp2 = (grad_ours==0.0).astype(float)
        assert np.sum(tmp1 - tmp2) == 0.0 
        assert torch.max(output_ref) == torch.max(output_ours)
        assert np.max(grad_ref) == np.max(grad_ours) 
    

def test_dropout_memory():
    print("========== Dropout Memory Test ==========")
    data_np = np.random.randn(1024, 128).astype('float32')

    def test_implementation(func, p, training):
        data = torch.tensor(data_np).to("cuda").requires_grad_()
        before = get_memory_usage(0)

        for i in range(10):
            data = func(data, p, training)

        after = get_memory_usage(0)

        return after - before
    p = 0.2
    training = True
    usage_ref = test_implementation(F.dropout, p, training)
    usage_us = test_implementation(ext_quantization.act_quantized_dropout, p, training)

    print("Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
    print("Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))


def test_spmm_matmul_correctness():
    print('============test spmm correctness============')
    N, D = 10000, 128
    nnz = 50000
    i = torch.randint(high=N, size=(2, nnz), dtype=torch.int64)
    v = torch.randn(size=(nnz,))
    dense_mat_cpu = torch.randn(N, D)
    def test_implementation_has_value(func, reduce):
        v_ = v.clone().requires_grad_()
        torch_sp = torch.sparse_coo_tensor(i, v_, [N, N])
        tsp = torch_sparse.SparseTensor.from_torch_sparse_coo_tensor(torch_sp).cuda()
        dense_mat = dense_mat_cpu.cuda().requires_grad_()
        output = func(tsp, dense_mat, reduce)
        output.backward(torch.ones_like(output))
        return [x.detach().cpu().numpy() for x in [output, dense_mat.grad, v_.grad]]

    def test_implementation_non_value(func, reduce):
        tsp = torch_sparse.SparseTensor.from_edge_index(i, sparse_sizes=[N, N]).cuda()
        dense_mat = dense_mat_cpu.cuda().requires_grad_()
        output = func(tsp, dense_mat, reduce)
        output.backward(torch.ones_like(output))
        return [x.detach().cpu().numpy() for x in [output, dense_mat.grad]]

    for reduce in ['sum', 'mean', 'max', 'min']:
        print(f'============test spmm {reduce} correctness============')
        output_ref, grad_data_ref, grad_value_ref = test_implementation_has_value(torch_sparse.matmul, reduce)
        output_ref_nonvalue, grad_data_ref_nonvalue  = test_implementation_non_value(torch_sparse.matmul, reduce)
        value_grads, data_grads = [], []
        for _ in range(10):
            output_us, grad_data_us, grad_value_us = test_implementation_has_value(qmatmul, reduce)
            np.testing.assert_allclose(output_ref, output_us)
            value_grads.append(grad_value_us)
            data_grads.append(grad_data_us)
        value_grads = np.stack(value_grads, 0)
        value_grads_mean = value_grads.mean(0)
        value_grads_std = value_grads.std(0)
        bias = np.linalg.norm(value_grads_mean - grad_value_ref)
        print('Value Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(grad_value_ref), bias, np.linalg.norm(value_grads_std)))

        data_grads = np.stack(data_grads, 0)
        data_grads_mean = data_grads.mean(0)
        data_grads_std = data_grads.std(0)
        bias = np.linalg.norm(data_grads_mean - grad_data_ref)
        print('Data Grad = {}, Bias = {}, Std = {}'.format(np.linalg.norm(grad_data_ref), bias, np.linalg.norm(data_grads_std)))

        data_grads = []
        for _ in range(10):
            output_us, grad_data_us = test_implementation_non_value(qmatmul, reduce)
            np.testing.assert_allclose(output_ref_nonvalue, output_us)
            data_grads.append(grad_data_us)

        data_grads = np.stack(data_grads, 0)
        data_grads_mean = data_grads.mean(0)
        data_grads_std = data_grads.std(0)
        bias = np.linalg.norm(data_grads_mean - grad_data_ref_nonvalue)
        print('(value.requires_grad = False) Data Grad = {}, Bias = {}, Std = {}'.format(
            np.linalg.norm(grad_data_ref_nonvalue), bias, np.linalg.norm(data_grads_std)))


def test_spmm_matmul_memory():
    print("========== Spmm Memory Test ==========")
    N, D = 10000, 128
    nnz = 50000
    data_np = np.random.randn(N, D).astype('float32')
    # w = torch.randn((D, D), requires_grad=True, device='cuda')
    # qw = torch.randn((D, D), requires_grad=True, device='cuda')
    # with torch.no_grad():
    #     qw.copy_(w)
    index = torch.randint(high=N, size=(2, nnz), dtype=torch.int64)
    value = torch.randn(size=(nnz,))
    torch_sp = torch.sparse_coo_tensor(index, value, [N, N])

    def test_implementation(func, n_layers, reduce):
        data = torch.tensor(data_np).to("cuda").requires_grad_()
        tsp = torch_sparse.SparseTensor.from_torch_sparse_coo_tensor(torch_sp).cuda()
        tsp.fill_cache_()
        output = data
        before = get_memory_usage(0)

        for i in range(n_layers):
            output = func(tsp, output, reduce)

        after = get_memory_usage(0) - compute_tensor_bytes([output])
        
        if func == torch_sparse.matmul:
            after += compute_tensor_bytes(data)
        return after - before
    
    for reduce in ['sum', 'mean', 'max', 'min']:
        print(f'============test spmm {reduce} Memory============')    
        usage_ref = test_implementation(torch_sparse.matmul, 10, reduce)
        usage_us = test_implementation(qmatmul, 10, reduce)
        print("Exact.     Usage: %.2f MB" % (usage_ref / 2 ** 20))
        print("Quantized. Usage: %.2f MB" % (usage_us / 2 ** 20))
        print("Ratio: %.2f" % (usage_ref / usage_us))
        print("")


if __name__ == "__main__":
    # test_relu_correctness()
    # test_relu_memory()
    # test_relu_speed()
    # test_leaky_relu_correctness()
    # test_leaky_relu_memory()
    # test_leaky_relu_speed()
    # test_elu_correctness()
    # test_elu_memory()
    # test_linear_correctness()
    # test_linear_memory()
    # test_linear_speed()
    # print(torch.backends.cudnn.version())
    # test_bn_correctness()
    # test_bn_speed()
    # test_dropout_correctness()
    # test_dropout_memory()
    # test_dropout_speed()
    test_spmm_matmul_correctness()
    # test_spmm_matmul_memory()
test/test_op.py
saintslab/i-Exact