<filename>examples/pytorch/graphsage/check-sparsity.py<gh_stars>0
"""
Inductive Representation Learning on Large Graphs
Paper: http://papers.nips.cc/paper/6703-inductive-representation-learning-on-large-graphs.pdf
Code: https://github.com/williamleif/graphsage-simple
Simple reference implementation of GraphSAGE.
"""
import argparse
import time
import numpy as np
import networkx as nx
import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
from dgl import DGLGraph
from dgl.data import register_data_args, load_data
from dgl.nn.pytorch.conv import SAGEConv
import sys
import matplotlib.pyplot as plt



class GraphSAGE(nn.Module):
    def __init__(self,
                 in_feats,
                 n_hidden,
                 n_classes,
                 n_layers,
                 activation,
                 dropout,
                 aggregator_type,
                 th1,th2,th3):
        super(GraphSAGE, self).__init__()
        self.layers = nn.ModuleList()
        self.dropout = nn.Dropout(dropout)
        self.activation = activation
        self.sparsify = False
        self.sparsify_threshold = th3
        self.prune_weight_threshold = th1
        self.prune_weight_threshold1 = th2 
        #print(th1, th2)
        #print("self.sparsify_threshold, self.prune_weight_threshold")
        #print(self.sparsify_threshold, self.prune_weight_threshold)


        # input layer
        self.layers.append(SAGEConv(in_feats, n_hidden, aggregator_type))
        # hidden layers
        for i in range(n_layers - 1):
            self.layers.append(SAGEConv(n_hidden, n_hidden, aggregator_type))
        # output layer
        self.layers.append(SAGEConv(n_hidden, n_classes, aggregator_type)) # activation None

    def forward(self, graph, inputs):
        h = self.dropout(inputs)
        

        for l, layer in enumerate(self.layers):
            if( self.sparsify == True):
                x = 0
                i,j = h.size()
                for in_data in h:
                     x = x + len(in_data[abs(in_data) > args.sparsify_feat_thd[l] ])
                sprs_rate = 1 - x/(j*i)
                print("\nLayer %d thd %4f feature sparse rate: %4.4f " % (l, args.sparsify_feat_thd[l], sprs_rate))
                check_input_sparsity(h)
                h[ abs(h)< args.sparsify_feat_thd[l]] = 0
                
            
            h = layer(graph, h)
            x, y = h.shape
            #print("x y   ", x, y)
            #print(h)
            if l != len(self.layers) - 1:
                fileName = args.dataset+"-layer"+str(l)+"-feat-sprs-mask.txt"
                f= open(fileName,"w+")
                i,j = h.size()
                feat_sparse_mask = [0]*j
                for in_data in h:
                     idx = 0
                     for element  in in_data:
                         
                         if(element >= args.sparsify_feat_thd[l]):
                             feat_sparse_mask[idx] = 1
                         idx = idx + 1

                     mask_output = ""
                     for ele in feat_sparse_mask:  
                        mask_output += str(ele) 
                     mask_output = mask_output + "\n"
                     f.write(mask_output)
                f.close()     

                h = self.activation(h)
                h = self.dropout(h)
        return h


def evaluate(model, graph, features, labels, nid):
    model.eval()
    with torch.no_grad():
        logits = model(graph, features)
        logits = logits[nid]
        labels = labels[nid]
        _, indices = torch.max(logits, dim=1)
        correct = torch.sum(indices == labels)
        return correct.item() * 1.0 / len(labels)

def examine_sparsity(model, graph, features, labels, nid):
    model.eval()
    with torch.no_grad():
        logits = model(graph, features)
        logits = logits[nid]
        labels = labels[nid]
        _, indices = torch.max(logits, dim=1)
        correct = torch.sum(indices == labels)
        return correct.item() * 1.0 / len(labels)


def prune( model, skip_last):
    pruned_weights = 0
    cnt  = 0
    layer_idx = 0
    #print("prune( ) threshold", model.prune_weight_threshold)
    for name, param in model.named_parameters():
        if param.requires_grad:
            #print(name, param.data)
            #print("shape: ",name, param.data.shape)
            
            if name.split('.')[-1] == "weight":
               layer_idx = layer_idx+1
               pruned_weights_layer = 0
               cnt_layer  = 0
               dim_out, dim_in = param.data.shape
               print("Out   In ", dim_out, dim_in)
               weights = torch.arange(dim_out*dim_in)
               weights2D = weights.view(dim_out, dim_in)

               if skip_last == True and dim_out < 16:
                   continue

               if layer_idx == 1:
                   thd = model.prune_weight_threshold
               else:
                   thd = model.prune_weight_threshold1

               _weight = abs(param.data)
               a, b = _weight.size()
               _weight = _weight.view(1, a*b)
               hist_w = np.histogram(_weight, bins='auto')
               print("hist_w"+args.dataset+"-"+name)
               print(hist_w)

               for i in range(dim_out):
                    for j in range(dim_in):
                        if abs(param.data[i,j]) < thd:
                            #print(param.data[i,j])
                            param.data[i,j] = 0
                            weights2D[i,j] = 0
                            pruned_weights = pruned_weights +1
                            pruned_weights_layer = pruned_weights_layer +1
                        else:
                            weights2D[i,j] = 1   
                        
                        cnt = cnt +1
                        cnt_layer = cnt_layer +1
                
               #print(name, " weights2D: ",dim_out, dim_in)
               #print(weights2D)
               fig = plt.figure()
               plt.imshow(weights2D)
               fig.suptitle(args.dataset+" "+name+" Thld "+str(thd)+" sp-rate "+str(pruned_weights_layer/cnt_layer))
               plt.savefig(args.dataset+"-"+name+".jpg")
               

               #weights2D = weights2D.view(dim_out*64,dim_in/64)
               step = int(dim_in/8)
               fig = plt.figure()
               
               #ax1 = fig.add_subplot(2,1,1)
               fig, axs = plt.subplots(8)
               for i in range(8):
                    axs[i].imshow(weights2D[:,i*step:(i+1)*step-1])
               fig.suptitle(args.dataset+" Weight Map with Pruning Weight Threshold  "+str(model.prune_weight_threshold))
               fig.savefig(args.dataset+"-"+name+"All.jpg")
               plt.close("all")
               print("prune rate ", pruned_weights, cnt, pruned_weights/cnt,layer_idx)
            
    print("prune rate ", pruned_weights, cnt, pruned_weights/cnt)


class Hook():
    def __init__(self, module, backward=False):
        if backward==False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)
    def hook_fn(self, module, input, output):
        self.input = input
        self.output = output
    def close(self):
        self.hook.remove()

def hook_fn_forward(module, input_feat, output_feat):
    #print("Hook ",module,input_feat[0].size(),output_feat.shape) 
    #print('input data:', input_feat[0]) 
    x = 0
    i,j = input_feat[0].size()
   # for in_data in input_feat[0]:
   #      x = x + len(in_data[abs(in_data) > args.sparsify_feature_threshold ])
   # print("\n Input feature spars rate: ", module, 1 - x/(j*i))
   # check_input_sparsity( input_feat[0])


def check_input_sparsity( features):
    x , y = features.size()
    print("check %d nodes %d features " %(x , y ))
    
    features_m = abs(features)
    
    node_set = range(1,x)
    feature_set = range(1, y)
    features_flat = torch.reshape(features_m,(1, x*y))   #features_m.resize(1, x*y)
    
    feat_sort, indices = torch.sort(features_flat,1 , True)
    f1 = feat_sort.data.numpy()
    feat_sum = f1.sum()
    x1, y1 = np.shape(f1)
    print(feat_sort)
    cmu_sum = 0
    num = x*y
    print(x1, y1, x, y, num, feat_sum,f1[0][0],f1[0][1],f1[0][2])
    i = 0
    
    thd = 0.85
    for val_set in f1:
      for val in val_set:
         i = i +1
         cmu_sum = cmu_sum + val
         #print(i, cmu_sum, val)
        
         if cmu_sum/feat_sum >= thd:
             print(" Percent. %3.1f value  %f  %f  %f  percent_nodes %4.4f" % (thd*100, val, cmu_sum, feat_sum, i/num))
             break
                    
    #for node in features_m:
    #  print(torch.mean(node), torch.max(node),torch.var(node),1-len(torch.nonzero(node))/len(node))

    


def main(args):
    # load and preprocess dataset
    data = load_data(args)
    g = data[0]
    features = g.ndata['feat']
    labels = g.ndata['label']
    train_mask = g.ndata['train_mask']
    val_mask = g.ndata['val_mask']
    test_mask = g.ndata['test_mask']
    in_feats = features.shape[1]
    n_classes = data.num_classes
    n_edges = data.graph.number_of_edges()
    print("""----Data statistics------'
      #Edges %d
      #Classes %d
      #Train samples %d
      #Val samples %d
      #Test samples %d""" %
          (n_edges, n_classes,
           train_mask.int().sum().item(),
           val_mask.int().sum().item(),
           test_mask.int().sum().item()))
    
    #check_input_sparsity( features)
    
    

    if args.gpu < 0:
        cuda = False
    else:
        cuda = True
        torch.cuda.set_device(args.gpu)
        features = features.cuda()
        labels = labels.cuda()
        train_mask = train_mask.cuda()
        val_mask = val_mask.cuda()
        test_mask = test_mask.cuda()
        print("use cuda:", args.gpu)

    train_nid = train_mask.nonzero().squeeze()
    val_nid = val_mask.nonzero().squeeze()
    test_nid = test_mask.nonzero().squeeze()

    # graph preprocess and calculate normalization factor
    g = dgl.remove_self_loop(g)
    n_edges = g.number_of_edges()
    if cuda:
        g = g.int().to(args.gpu)

    # create GraphSAGE model
    model = GraphSAGE(in_feats,
                      args.n_hidden,
                      n_classes,
                      args.n_layers,
                      F.relu,
                      args.dropout,
                      args.aggregator_type,
                      args.prune_weight_threshold,
                      args.prune_weight_threshold1,
                      args.sparsify_feature_threshold)

    if cuda:
        model.cuda()
    
    #print("Show in_feat:")
    #print(in_feats)
    #print("show  features")
    #print(features)

    #sys.exit("Done init")

    # use optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    
    check_sparsity = True
    
       
    # initialize graph
    dur = []
    for epoch in range(args.n_epochs):
        
        if check_sparsity == True:
                input_graph = args.dataset+"-ckpt"
                ckpt = torch.load(input_graph)
                print( ckpt)
                model.load_state_dict(ckpt['model_state_dict'])
                break

        model.train()
        if epoch >= 3:
            t0 = time.time()
                
                

        # forward
        logits = model(g, features)
        loss = F.cross_entropy(logits[train_nid], labels[train_nid])

        optimizer.zero_grad()
        loss.backward()

        #prune( model)

        optimizer.step()

        if epoch >= 3:
            dur.append(time.time() - t0)

        acc = evaluate(model, g, features, labels, val_nid)
        print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
              "ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
                                            acc, n_edges / np.mean(dur) / 1000))

    
    #print("modules: ")
    #for name, module in model.named_modules():
    #    print(name)

    
    acc = evaluate(model, g, features, labels, test_nid)
    print("Test Accuracy {:.4f}".format(acc))

    #print("Show in_feat:")
    #print(in_feats)
    #print("show  features")
    #print(features)
    #print("show  model")
    #print(model)
    #for name, param in model.named_parameters():
    #    if param.requires_grad:
    #        print(name, param.data)
    #        print("shape: ",param.data.shape)
            
    #print()
    #print("show list(model._modules.items() ")
    #for m in list(model._modules.items()):
    #for m in model.layers:
    #    print(m)
        

    

    #hookF = [Hook(layer[1]) for layer in list(model._modules.items())]
    #acc = evaluate(model, g, features, labels, test_nid)
    #print('***'*3+'  Forward Hooks Inputs & Outputs  '+'***'*3)
    #for hook in hookF:
    #    print(hook.input)
    #    print(hook.output)
    #    print('---'*17)
    #print('\n')
   
    
    modules = model.named_children() # 
    for name, module in modules:
       print("register a hook for:", name)
       module.register_forward_hook(hook_fn_forward)

    #print("show features and labels")
    #print( features.size(), labels.size( ))

    #print(features)
    #print(labels)   
    
    prune_model = False
    sparse_input = True

    if prune_model == True:
        prune(model,True)
        acc = evaluate(model, g, features, labels, test_nid)
        print("prune-weight-threshold:  ",model.prune_weight_threshold)
        print("prune-weight-threshold1:  ",model.prune_weight_threshold1)
        print("Test Accuracy with Pruning(skip) {:.4f}".format(acc)) 
    
        torch.save({
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict()
            
            }, args.dataset+"-Sparse-ckpt")

    if sparse_input == True: 
        model.sparsify = True
        acc = evaluate(model, g, features, labels, test_nid)
        print("\nTest Accuracy with Sparsify  {:.4f}".format(acc))  
    
    return


    #print()
    prune(model,False)
    acc = evaluate(model, g, features, labels, test_nid)
    print("Test Accuracy for Pruned with False {:.4f}".format(acc))        

    torch.save(model.state_dict(), "my_model")

    


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='GraphSAGE')
    register_data_args(parser)
    parser.add_argument("--dropout", type=float, default=0.5,
                        help="dropout probability")
    parser.add_argument("--gpu", type=int, default=-1,
                        help="gpu")
    parser.add_argument("--lr", type=float, default=1e-2,
                        help="learning rate")
    parser.add_argument("--n-epochs", type=int, default=200,
                        help="number of training epochs")
    parser.add_argument("--n-hidden", type=int, default=16,
                        help="number of hidden gcn units")
    parser.add_argument("--n-layers", type=int, default=1,
                        help="number of hidden gcn layers")
    parser.add_argument("--weight-decay", type=float, default=5e-4,
                        help="Weight for L2 loss")
    parser.add_argument("--prune-weight-threshold", type=float, default=0.125,
                        help="prune-weight-threshold") 
    parser.add_argument("--prune-weight-threshold1", type=float, default=0.01,
                        help="prune-weight-threshold1")                     
    parser.add_argument("--sparsify-feature-threshold", type=float, default=0.01,
                        help="sparsify-feature-threshold")                                       
    parser.add_argument("--aggregator-type", type=str, default="gcn",
                        help="Aggregator type: mean/gcn/pool/lstm")
    parser.add_argument('--sparsify-feat-thd', nargs='+', type=float, help="layer-wise sparsify-feature-threshold")
    args = parser.parse_args()
    print(args)

    main(args)
