First commit

README.md

@@ -1,8 +1,8 @@
 ---
-title: AioMedica2
-emoji: 👁
-colorFrom: red
-colorTo: indigo
+title: AioMedica
+emoji: 🏃
+colorFrom: purple
+colorTo: yellow
 sdk: streamlit
 sdk_version: 1.10.0
 app_file: app.py

app.py

@@ -0,0 +1,175 @@
+import streamlit as st
+import openslide
+import os
+from streamlit_option_menu import option_menu
+import torch
+
+
+@st.cache(suppress_st_warning=True)
+def load_model():
+    from predict import Predictor
+    predictor = Predictor()
+    return predictor
+
+@st.cache(suppress_st_warning=True)
+def load_dependencies():
+        
+    os.system("pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.1+cpu.html")
+    os.system("pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.7.1+cpu.html")
+    os.system("pip install torch-geometric -f https://pytorch-geometric.com/whl/torch-1.7.1+cpu.html")
+
+
+def main():
+
+        
+
+
+
+        # environment variables for the inference api
+        os.environ['DATA_DIR'] = 'queries'
+        os.environ['PATCHES_DIR'] =  os.path.join(os.environ['DATA_DIR'], 'patches')
+        os.environ['SLIDES_DIR'] = os.path.join(os.environ['DATA_DIR'], 'slides')
+        os.environ['GRAPHCAM_DIR'] = os.path.join(os.environ['DATA_DIR'], 'graphcam_plots')
+        os.makedirs(os.environ['GRAPHCAM_DIR'], exist_ok=True)
+
+
+        # manually put the metadata in the metadata folder
+        os.environ['CLASS_METADATA'] ='metadata/label_map.pkl'
+
+        # manually put the desired weights in the weights folder
+        os.environ['WEIGHTS_PATH'] = WEIGHTS_PATH='weights'
+        os.environ['FEATURE_EXTRACTOR_WEIGHT_PATH'] = os.path.join(os.environ['WEIGHTS_PATH'], 'feature_extractor', 'model.pth')
+        os.environ['GT_WEIGHT_PATH']  = os.path.join(os.environ['WEIGHTS_PATH'], 'graph_transformer', 'GraphCAM.pth')
+
+
+        #st.set_page_config(page_title="",layout='wide')
+        predictor = load_model()#Predictor()
+
+
+
+
+
+        ABOUT_TEXT = "🤗 LastMinute Medical - Web diagnosis tool."
+        CONTACT_TEXT = """
+        _Built by Christian Cancedda and LabLab lads with love_ ❤️ 
+        [![Follow](https://img.shields.io/github/followers/Chris1nexus?style=social)](https://github.com/Chris1nexus)
+        [![Follow](https://img.shields.io/twitter/follow/chris_cancedda?style=social)](https://twitter.com/intent/follow?screen_name=chris_cancedda)
+        Star project repository:
+        [![GitHub stars](https://img.shields.io/github/followers/Chris1nexus?style=social)](https://github.com/Chris1nexus/inference-graph-transformer)
+        """
+        VISUALIZE_TEXT = "Visualize WSI slide by uploading it on the provided window"
+        DETECT_TEXT = "Generate a preliminary diagnosis about the presence of  pulmonary disease"
+
+
+
+        with st.sidebar:
+            choice = option_menu("LastMinute - Diagnosis",
+                                 ["About", "Visualize WSI slide", "Cancer Detection", "Contact"],
+                                 icons=['house', 'upload', 'activity',  'person lines fill'],
+                                 menu_icon="app-indicator", default_index=0,
+                                 styles={
+                                     # "container": {"padding": "5!important", "background-color": "#fafafa", },
+                                     "container": {"border-radius": ".0rem"},
+                                     # "icon": {"color": "orange", "font-size": "25px"},
+                                     # "nav-link": {"font-size": "16px", "text-align": "left", "margin": "0px",
+                                     #              "--hover-color": "#eee"},
+                                     # "nav-link-selected": {"background-color": "#02ab21"},
+                                 }
+                                 )
+        st.sidebar.markdown(
+            """
+        <style>
+        .aligncenter {
+            text-align: center;
+        }
+        </style>
+        <p style='text-align: center'>
+        <a href="https://github.com/Chris1nexus/inference-graph-transformer" target="_blank">Project Repository</a>
+        </p>
+        <p class="aligncenter">
+
+            <a href="https://github.com/Chris1nexus/inference-graph-transformer" target="_blank"> 
+                <img src="https://img.shields.io/github/stars/Chris1nexus/inference-graph-transformer?style=social"/>
+            </a>
+        </p>
+
+        <p class="aligncenter">
+            <a href="https://twitter.com/chris_cancedda" target="_blank"> 
+                <img src="https://img.shields.io/twitter/follow/chris_cancedda?style=social"/>
+            </a>
+        </p>
+            """,
+            unsafe_allow_html=True,
+        )
+
+
+        if choice == "About":
+            st.title(choice)
+            README = requests.get("https://raw.githubusercontent.com/Chris1nexus/inference-graph-transformer/master/README.md").text
+            README = str(README).replace('width="1200"','width="700"')
+            # st.title(choose)
+            st.markdown(README, unsafe_allow_html=True)
+
+        if choice == "Visualize WSI slide":
+            st.title(choice)
+            st.markdown(VISUALIZE_TEXT)
+
+            uploaded_file = st.file_uploader("Choose a WSI slide file to diagnose (.svs)")
+            if uploaded_file is not None:
+                ori = openslide.OpenSlide(uploaded_file.name)
+                width, height = ori.dimensions
+
+                REDUCTION_FACTOR = 20
+                w, h = int(width/512), int(height/512)
+                w_r, h_r = int(width/20), int(height/20)
+                resized_img = ori.get_thumbnail((w_r,h_r))
+                resized_img = resized_img.resize((w_r,h_r))
+                ratio_w, ratio_h = width/resized_img.width, height/resized_img.height
+                #print('ratios ', ratio_w, ratio_h)
+                w_s, h_s = float(512/REDUCTION_FACTOR), float(512/REDUCTION_FACTOR)   
+                st.image(resized_img, use_column_width='never')   
+
+        if choice == "Cancer Detection":
+            state = dict()
+
+            st.title(choice)
+            st.markdown(DETECT_TEXT)
+            uploaded_file = st.file_uploader("Choose a WSI slide file to diagnose (.svs)")
+            st.markdown("Examples can be chosen at the [GDC Data repository](https://portal.gdc.cancer.gov/repository?facetTab=cases&filters=%7B%22op%22%3A%22and%22%2C%22content%22%3A%5B%7B%22op%22%3A%22in%22%2C%22content%22%3A%7B%22field%22%3A%22cases.primary_site%22%2C%22value%22%3A%5B%22bronchus%20and%20lung%22%5D%7D%7D%2C%7B%22op%22%3A%22in%22%2C%22content%22%3A%7B%22field%22%3A%22cases.project.program.name%22%2C%22value%22%3A%5B%22TCGA%22%5D%7D%7D%2C%7B%22op%22%3A%22in%22%2C%22content%22%3A%7B%22field%22%3A%22cases.project.project_id%22%2C%22value%22%3A%5B%22TCGA-LUAD%22%2C%22TCGA-LUSC%22%5D%7D%7D%2C%7B%22op%22%3A%22in%22%2C%22content%22%3A%7B%22field%22%3A%22files.experimental_strategy%22%2C%22value%22%3A%5B%22Tissue%20Slide%22%5D%7D%7D%5D%7D)")
+            st.markdown("Alternatively, for simplicity few test cases are provided at the [drive link](https://drive.google.com/drive/folders/1u3SQa2dytZBHHh6eXTlMKY-pZGZ-pwkk?usp=share_link)")
+
+
+            if uploaded_file is not None:
+                # To read file as bytes:
+                #print(uploaded_file)
+                with open(os.path.join(uploaded_file.name),"wb") as f:
+                     f.write(uploaded_file.getbuffer())
+                with st.spinner(text="Computation is running"):
+                    predicted_class, viz_dict = predictor.predict(uploaded_file.name)
+                st.info('Computation completed.')
+                st.header(f'Predicted to be: {predicted_class}')
+                st.text('Heatmap of the areas that show markers correlated with the disease.\nIncreasing red tones represent higher likelihood that the area is affected')
+                state['cur'] = predicted_class
+                mapper = {'ORI': predicted_class, predicted_class:'ORI'}
+                readable_mapper = {'ORI': 'Original',  predicted_class :'Disease heatmap' }        
+                #def fn():
+                #    st.image(viz_dict[mapper[state['cur']]], use_column_width='never', channels='BGR') 
+                #    state['cur'] = mapper[state['cur']]
+                #    return 
+
+                #st.button(f'See {readable_mapper[mapper[state["cur"]] ]}', on_click=fn   )
+                #st.image(viz_dict[state['cur']], use_column_width='never', channels='BGR') 
+                st.image([viz_dict[state['cur']],viz_dict['ORI']], caption=['Original', f'{predicted_class} heatmap'] ,channels='BGR'
+                    # use_column_width='never', 
+                    ) 
+                    
+
+        if choice == "Contact":
+            st.title(choice)
+            st.markdown(CONTACT_TEXT)
+
+if __name__ == '__main__':
+    #''' 
+    load_dependencies()
+    #'''
+    main()

feature_extractor/build_graph_utils.py

@@ -0,0 +1,85 @@
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import torchvision.models as models
+import torchvision.transforms.functional as VF
+from torchvision import transforms
+
+import sys, argparse, os, glob
+import pandas as pd
+import numpy as np
+from PIL import Image
+from collections import OrderedDict
+
+class ToPIL(object):
+    def __call__(self, sample):
+        img = sample
+        img = transforms.functional.to_pil_image(img)
+        return img 
+
+class BagDataset():
+    def __init__(self, csv_file, transform=None):
+        self.files_list = csv_file
+        self.transform = transform
+    def __len__(self):
+        return len(self.files_list)
+    def __getitem__(self, idx):
+        temp_path = self.files_list[idx]
+        img = os.path.join(temp_path)
+        img = Image.open(img)
+        img = img.resize((224, 224))
+        sample = {'input': img}
+        
+        if self.transform:
+            sample = self.transform(sample)
+        return sample 
+
+class ToTensor(object):
+    def __call__(self, sample):
+        img = sample['input']
+        img = VF.to_tensor(img)
+        return {'input': img} 
+    
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, img):
+        for t in self.transforms:
+            img = t(img)
+        return img
+
+def save_coords(txt_file, csv_file_path):
+    for path in csv_file_path:
+        x, y = path.split('/')[-1].split('.')[0].split('_')
+        txt_file.writelines(str(x) + '\t' + str(y) + '\n')
+    txt_file.close()
+
+def adj_matrix(csv_file_path, output,  device='cpu'):
+    total = len(csv_file_path)
+    adj_s = np.zeros((total, total))
+
+    for i in range(total-1):
+        path_i = csv_file_path[i]
+        x_i, y_i = path_i.split('/')[-1].split('.')[0].split('_')
+        for j in range(i+1, total):
+            # sptial 
+            path_j = csv_file_path[j]
+            x_j, y_j = path_j.split('/')[-1].split('.')[0].split('_')
+            if abs(int(x_i)-int(x_j)) <=1 and abs(int(y_i)-int(y_j)) <= 1:
+                adj_s[i][j] = 1
+                adj_s[j][i] = 1
+
+    adj_s = torch.from_numpy(adj_s)
+    adj_s = adj_s.to(device)
+
+    return adj_s
+
+def bag_dataset(args, csv_file_path):
+    transformed_dataset = BagDataset(csv_file=csv_file_path,
+                                    transform=Compose([
+                                        ToTensor()
+                                    ]))
+    dataloader = DataLoader(transformed_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, drop_last=False)
+    return dataloader, len(transformed_dataset)

feature_extractor/build_graphs.py

@@ -0,0 +1,114 @@
+
+from cl import IClassifier 
+from build_graph_utils import *
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import torchvision.models as models
+import torchvision.transforms.functional as VF
+from torchvision import transforms
+
+import sys, argparse, os, glob
+import pandas as pd
+import numpy as np
+from PIL import Image
+from collections import OrderedDict
+
+
+
+def compute_feats(args, bags_list, i_classifier, device, save_path=None, whole_slide_path=None):
+    num_bags = len(bags_list)
+    Tensor = torch.FloatTensor
+    for i in range(0, num_bags):
+        feats_list = []
+        if  args.magnification == '20x':
+            glob_path = os.path.join(bags_list[i], '*.jpeg')
+            csv_file_path = glob.glob(glob_path)
+            # line below was in the original version, commented due to errror with current version
+            #file_name = bags_list[i].split('/')[-3].split('_')[0]
+            
+            file_name = glob_path.split('/')[-3].split('_')[0]
+            
+        if args.magnification == '5x' or args.magnification == '10x':
+            csv_file_path = glob.glob(os.path.join(bags_list[i], '*.jpg'))
+
+        dataloader, bag_size = bag_dataset(args, csv_file_path)
+        print('{} files to be processed: {}'.format(len(csv_file_path), file_name))
+
+        if os.path.isdir(os.path.join(save_path, 'simclr_files', file_name)) or len(csv_file_path) < 1:
+            print('alreday exists')
+            continue
+        with torch.no_grad():
+            for iteration, batch in enumerate(dataloader):
+                patches = batch['input'].float().to(device) 
+                feats, classes = i_classifier(patches)
+                #feats = feats.cpu().numpy()
+                feats_list.extend(feats)
+        
+        os.makedirs(os.path.join(save_path, 'simclr_files', file_name), exist_ok=True)
+
+        txt_file = open(os.path.join(save_path, 'simclr_files', file_name, 'c_idx.txt'), "w+")
+        save_coords(txt_file, csv_file_path)
+        # save node features
+        output = torch.stack(feats_list, dim=0).to(device)
+        torch.save(output, os.path.join(save_path, 'simclr_files', file_name, 'features.pt'))
+        # save adjacent matrix
+        adj_s = adj_matrix(csv_file_path, output, device=device)
+        torch.save(adj_s, os.path.join(save_path, 'simclr_files', file_name, 'adj_s.pt'))
+
+        print('\r Computed: {}/{}'.format(i+1, num_bags))
+        
+
+def main():
+    parser = argparse.ArgumentParser(description='Compute TCGA features from SimCLR embedder')
+    parser.add_argument('--num_classes', default=2, type=int, help='Number of output classes')
+    parser.add_argument('--num_feats', default=512, type=int, help='Feature size')
+    parser.add_argument('--batch_size', default=128, type=int, help='Batch size of dataloader')
+    parser.add_argument('--num_workers', default=0, type=int, help='Number of threads for datalodaer')
+    parser.add_argument('--dataset', default=None, type=str, help='path to patches')
+    parser.add_argument('--backbone', default='resnet18', type=str, help='Embedder backbone')
+    parser.add_argument('--magnification', default='20x', type=str, help='Magnification to compute features')
+    parser.add_argument('--weights', default=None, type=str, help='path to the pretrained weights')
+    parser.add_argument('--output', default=None, type=str, help='path to the output graph folder')
+    args = parser.parse_args()
+    
+    if args.backbone == 'resnet18':
+        resnet = models.resnet18(pretrained=False, norm_layer=nn.InstanceNorm2d)
+        num_feats = 512
+    if args.backbone == 'resnet34':
+        resnet = models.resnet34(pretrained=False, norm_layer=nn.InstanceNorm2d)
+        num_feats = 512
+    if args.backbone == 'resnet50':
+        resnet = models.resnet50(pretrained=False, norm_layer=nn.InstanceNorm2d)
+        num_feats = 2048
+    if args.backbone == 'resnet101':
+        resnet = models.resnet101(pretrained=False, norm_layer=nn.InstanceNorm2d)
+        num_feats = 2048
+    for param in resnet.parameters():
+        param.requires_grad = False
+    resnet.fc = nn.Identity()
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print("Running on:", device)
+    i_classifier = IClassifier(resnet, num_feats, output_class=args.num_classes).to(device)
+    
+    # load feature extractor
+    if args.weights is None:
+        print('No feature extractor')
+        return
+    state_dict_weights = torch.load(args.weights)
+    state_dict_init = i_classifier.state_dict()
+    new_state_dict = OrderedDict()
+    for (k, v), (k_0, v_0) in zip(state_dict_weights.items(), state_dict_init.items()):
+        if 'features' not in k:
+            continue        
+        name = k_0
+        new_state_dict[name] = v
+    i_classifier.load_state_dict(new_state_dict, strict=False)
+ 
+    os.makedirs(args.output, exist_ok=True)
+    bags_list = glob.glob(args.dataset)
+    print(bags_list)
+    compute_feats(args, bags_list, i_classifier, device, args.output)
+    
+if __name__ == '__main__':
+    main()

feature_extractor/cl.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+class FCLayer(nn.Module):
+    def __init__(self, in_size, out_size=1):
+        super(FCLayer, self).__init__()
+        self.fc = nn.Sequential(nn.Linear(in_size, out_size))
+    def forward(self, feats):
+        x = self.fc(feats)
+        return feats, x
+
+class IClassifier(nn.Module):
+    def __init__(self, feature_extractor, feature_size, output_class):
+        super(IClassifier, self).__init__()
+        
+        self.feature_extractor = feature_extractor      
+        self.fc = nn.Linear(feature_size, output_class)
+        
+        
+    def forward(self, x):
+        device = x.device
+        feats = self.feature_extractor(x) # N x K
+        c = self.fc(feats.view(feats.shape[0], -1)) # N x C
+        return feats.view(feats.shape[0], -1), c
+
+class BClassifier(nn.Module):
+    def __init__(self, input_size, output_class, dropout_v=0.0): # K, L, N
+        super(BClassifier, self).__init__()
+        self.q = nn.Linear(input_size, 128)
+        self.v = nn.Sequential(
+            nn.Dropout(dropout_v),
+            nn.Linear(input_size, input_size)
+        )
+        
+        ### 1D convolutional layer that can handle multiple class (including binary)
+        self.fcc = nn.Conv1d(output_class, output_class, kernel_size=input_size)  
+        
+    def forward(self, feats, c): # N x K, N x C
+        device = feats.device
+        V = self.v(feats) # N x V, unsorted
+        Q = self.q(feats).view(feats.shape[0], -1) # N x Q, unsorted
+        
+        # handle multiple classes without for loop
+        _, m_indices = torch.sort(c, 0, descending=True) # sort class scores along the instance dimension, m_indices in shape N x C
+        m_feats = torch.index_select(feats, dim=0, index=m_indices[0, :]) # select critical instances, m_feats in shape C x K 
+        q_max = self.q(m_feats) # compute queries of critical instances, q_max in shape C x Q
+        A = torch.mm(Q, q_max.transpose(0, 1)) # compute inner product of Q to each entry of q_max, A in shape N x C, each column contains unnormalized attention scores
+        A = F.softmax( A / torch.sqrt(torch.tensor(Q.shape[1], dtype=torch.float32, device=device)), 0) # normalize attention scores, A in shape N x C, 
+        B = torch.mm(A.transpose(0, 1), V) # compute bag representation, B in shape C x V
+        
+        
+#         for i in range(c.shape[1]):
+#             _, indices = torch.sort(c[:, i], 0, True)         
+#             feats = torch.index_select(feats, 0, indices) # N x K, sorted
+#             q_max = self.q(feats[0].view(1, -1)) # 1 x 1 x Q
+#             temp = torch.mm(Q, q_max.view(-1, 1)) / torch.sqrt(torch.tensor(Q.shape[1], dtype=torch.float32, device=device))
+#             if i == 0:
+#                 A = F.softmax(temp, 0) # N x 1
+#                 B = torch.sum(torch.mul(A, V), 0).view(1, -1) # 1 x V
+#             else:
+#                 temp = F.softmax(temp, 0) # N x 1
+#                 A = torch.cat((A, temp), 1) # N x C
+#                 B = torch.cat((B, torch.sum(torch.mul(temp, V), 0).view(1, -1)), 0) # C x V -> 1 x C x V
+                
+        B = B.view(1, B.shape[0], B.shape[1]) # 1 x C x V
+        C = self.fcc(B) # 1 x C x 1
+        C = C.view(1, -1)
+        return C, A, B 
+    
+class MILNet(nn.Module):
+    def __init__(self, i_classifier, b_classifier):
+        super(MILNet, self).__init__()
+        self.i_classifier = i_classifier
+        self.b_classifier = b_classifier
+        
+    def forward(self, x):
+        feats, classes = self.i_classifier(x)
+        prediction_bag, A, B = self.b_classifier(feats, classes)
+        
+        return classes, prediction_bag, A, B
+        

feature_extractor/config.yaml

@@ -0,0 +1,23 @@
+batch_size: 256
+epochs: 20
+eval_every_n_epochs: 1
+fine_tune_from: ''
+log_every_n_steps: 25
+weight_decay: 10e-6
+fp16_precision: False
+n_gpu: 2
+gpu_ids: (0,1)
+
+model:
+  out_dim: 512
+  base_model: "resnet18"
+
+dataset:
+  s: 1
+  input_shape: (224,224,3)
+  num_workers: 10
+  valid_size: 0.1
+
+loss:
+  temperature: 0.5
+  use_cosine_similarity: True

feature_extractor/data_aug/dataset_wrapper.py

@@ -0,0 +1,93 @@
+import numpy as np
+from torch.utils.data import DataLoader
+from torch.utils.data.sampler import SubsetRandomSampler
+import torchvision.transforms as transforms
+from data_aug.gaussian_blur import GaussianBlur
+from torchvision import datasets
+import pandas as pd
+from PIL import Image
+from skimage import io, img_as_ubyte
+
+np.random.seed(0)
+
+class Dataset():
+    def __init__(self, csv_file, transform=None):
+        lines = []
+        with open(csv_file) as f:
+            for line in f:
+                line = line.rstrip().strip()
+                lines.append(line)
+        self.files_list = lines#pd.read_csv(csv_file)
+        self.transform = transform
+    def __len__(self):
+        return len(self.files_list)
+    def __getitem__(self, idx):
+        temp_path = self.files_list[idx]# self.files_list.iloc[idx, 0]
+        img = Image.open(temp_path)
+        img = transforms.functional.to_tensor(img)
+        if self.transform:
+            sample = self.transform(img)
+        return sample
+
+class ToPIL(object):
+    def __call__(self, sample):
+        img = sample
+        img = transforms.functional.to_pil_image(img)
+        return img 
+
+class DataSetWrapper(object):
+
+    def __init__(self, batch_size, num_workers, valid_size, input_shape, s):
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.valid_size = valid_size
+        self.s = s
+        self.input_shape = eval(input_shape)
+
+    def get_data_loaders(self):
+        data_augment = self._get_simclr_pipeline_transform()
+        train_dataset = Dataset(csv_file='all_patches.csv', transform=SimCLRDataTransform(data_augment))
+        train_loader, valid_loader = self.get_train_validation_data_loaders(train_dataset)
+        return train_loader, valid_loader
+
+    def _get_simclr_pipeline_transform(self):
+        # get a set of data augmentation transformations as described in the SimCLR paper.
+        color_jitter = transforms.ColorJitter(0.8 * self.s, 0.8 * self.s, 0.8 * self.s, 0.2 * self.s)
+        data_transforms = transforms.Compose([ToPIL(),
+                                            #   transforms.RandomResizedCrop(size=self.input_shape[0]),
+                                              transforms.Resize((self.input_shape[0],self.input_shape[1])),
+                                              transforms.RandomHorizontalFlip(),
+                                              transforms.RandomApply([color_jitter], p=0.8),
+                                              transforms.RandomGrayscale(p=0.2),
+                                              GaussianBlur(kernel_size=int(0.06 * self.input_shape[0])),
+                                              transforms.ToTensor()])
+        return data_transforms
+
+    def get_train_validation_data_loaders(self, train_dataset):
+        # obtain training indices that will be used for validation
+        num_train = len(train_dataset)
+        indices = list(range(num_train))
+        np.random.shuffle(indices)
+
+        split = int(np.floor(self.valid_size * num_train))
+        train_idx, valid_idx = indices[split:], indices[:split]
+
+        # define samplers for obtaining training and validation batches
+        train_sampler = SubsetRandomSampler(train_idx)
+        valid_sampler = SubsetRandomSampler(valid_idx)
+
+        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, sampler=train_sampler,
+                                  num_workers=self.num_workers, drop_last=True, shuffle=False)
+        valid_loader = DataLoader(train_dataset, batch_size=self.batch_size, sampler=valid_sampler,
+                                  num_workers=self.num_workers, drop_last=True)
+        return train_loader, valid_loader
+
+
+class SimCLRDataTransform(object):
+    def __init__(self, transform):
+        self.transform = transform
+
+    def __call__(self, sample):
+        xi = self.transform(sample)
+        xj = self.transform(sample)
+        return xi, xj

feature_extractor/data_aug/gaussian_blur.py

@@ -0,0 +1,26 @@
+import cv2
+import numpy as np
+
+np.random.seed(0)
+
+
+class GaussianBlur(object):
+    # Implements Gaussian blur as described in the SimCLR paper
+    def __init__(self, kernel_size, min=0.1, max=2.0):
+        self.min = min
+        self.max = max
+        # kernel size is set to be 10% of the image height/width
+        self.kernel_size = kernel_size
+
+    def __call__(self, sample):
+        sample = np.array(sample)
+
+        # blur the image with a 50% chance
+        prob = np.random.random_sample()
+
+        if prob < 0.5:
+#            print(self.kernel_size)
+            sigma = (self.max - self.min) * np.random.random_sample() + self.min
+            sample = cv2.GaussianBlur(sample, (self.kernel_size, self.kernel_size), sigma)
+
+        return sample

feature_extractor/load_patches.py

@@ -0,0 +1,37 @@
+
+import os, glob
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data_path', type=str)
+    args = parser.parse_args()
+
+    wsi_slides_paths = []
+
+
+    def r(dirpath):
+        for file in os.listdir(dirpath):
+            path = os.path.join(dirpath, file)
+            if os.path.isfile(path) and file.endswith(".svs"):
+                wsi_slides_paths.append(path)
+            elif os.path.isdir(path):
+                r(path)
+    def r(dirpath):
+        for path in glob.glob(os.path.join(dirpath, '*','*.svs') ):#os.listdir(dirpath):
+            if os.path.isfile(path):
+                wsi_slides_paths.append(path)    
+    def r(dirpath):
+        for path in glob.glob(os.path.join(dirpath, '*', '*', '*.jpeg') ):#os.listdir(dirpath):
+            if os.path.isfile(path):
+                wsi_slides_paths.append(path)                             
+    r(args.data_path)
+    with open('all_patches.csv', 'w') as f:
+        for filepath in wsi_slides_paths:
+            f.write(f'{filepath}\n')
+
+
+
+
+if __name__ == "__main__":
+    main()

feature_extractor/loss/nt_xent.py

@@ -0,0 +1,65 @@
+import torch
+import numpy as np
+
+
+class NTXentLoss(torch.nn.Module):
+
+    def __init__(self, device, batch_size, temperature, use_cosine_similarity):
+        super(NTXentLoss, self).__init__()
+        self.batch_size = batch_size
+        self.temperature = temperature
+        self.device = device
+        self.softmax = torch.nn.Softmax(dim=-1)
+        self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool)
+        self.similarity_function = self._get_similarity_function(use_cosine_similarity)
+        self.criterion = torch.nn.CrossEntropyLoss(reduction="sum")
+
+    def _get_similarity_function(self, use_cosine_similarity):
+        if use_cosine_similarity:
+            self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
+            return self._cosine_simililarity
+        else:
+            return self._dot_simililarity
+
+    def _get_correlated_mask(self):
+        diag = np.eye(2 * self.batch_size)
+        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
+        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
+        mask = torch.from_numpy((diag + l1 + l2))
+        mask = (1 - mask).type(torch.bool)
+        return mask.to(self.device)
+
+    @staticmethod
+    def _dot_simililarity(x, y):
+        v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2)
+        # x shape: (N, 1, C)
+        # y shape: (1, C, 2N)
+        # v shape: (N, 2N)
+        return v
+
+    def _cosine_simililarity(self, x, y):
+        # x shape: (N, 1, C)
+        # y shape: (1, 2N, C)
+        # v shape: (N, 2N)
+        v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0))
+        return v
+
+    def forward(self, zis, zjs):
+        representations = torch.cat([zjs, zis], dim=0)
+
+        similarity_matrix = self.similarity_function(representations, representations)
+
+        # filter out the scores from the positive samples
+        l_pos = torch.diag(similarity_matrix, self.batch_size)
+        r_pos = torch.diag(similarity_matrix, -self.batch_size)
+        positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1)
+
+        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1)
+
+        logits = torch.cat((positives, negatives), dim=1)
+        logits /= self.temperature
+
+        labels = torch.zeros(2 * self.batch_size).to(self.device).long()
+        loss = self.criterion(logits, labels)
+
+        return loss / (2 * self.batch_size)

feature_extractor/models/baseline_encoder.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+
+class Encoder(nn.Module):
+    def __init__(self, out_dim=64):
+        super(Encoder, self).__init__()
+        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.conv4 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+
+        # projection MLP
+        self.l1 = nn.Linear(64, 64)
+        self.l2 = nn.Linear(64, out_dim)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.pool(x)
+
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = self.pool(x)
+
+        x = self.conv3(x)
+        x = F.relu(x)
+        x = self.pool(x)
+
+        x = self.conv4(x)
+        x = F.relu(x)
+        x = self.pool(x)
+
+        h = torch.mean(x, dim=[2, 3])
+
+        x = self.l1(h)
+        x = F.relu(x)
+        x = self.l2(x)
+
+        return h, x

feature_extractor/models/resnet_simclr.py

@@ -0,0 +1,37 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+
+class ResNetSimCLR(nn.Module):
+
+    def __init__(self, base_model, out_dim):
+        super(ResNetSimCLR, self).__init__()
+        self.resnet_dict = {"resnet18": models.resnet18(pretrained=False, norm_layer=nn.InstanceNorm2d),
+                            "resnet50": models.resnet50(pretrained=False)}
+
+        resnet = self._get_basemodel(base_model)
+        num_ftrs = resnet.fc.in_features
+
+        self.features = nn.Sequential(*list(resnet.children())[:-1])
+
+        # projection MLP
+        self.l1 = nn.Linear(num_ftrs, num_ftrs)
+        self.l2 = nn.Linear(num_ftrs, out_dim)
+
+    def _get_basemodel(self, model_name):
+        try:
+            model = self.resnet_dict[model_name]
+            print("Feature extractor:", model_name)
+            return model
+        except:
+            raise ("Invalid model name. Check the config file and pass one of: resnet18 or resnet50")
+
+    def forward(self, x):
+        h = self.features(x)
+        h = h.squeeze()
+
+        x = self.l1(h)
+        x = F.relu(x)
+        x = self.l2(x)
+        return h, x

feature_extractor/run.py

@@ -0,0 +1,21 @@
+from simclr import SimCLR
+import yaml
+from data_aug.dataset_wrapper import DataSetWrapper
+import os, glob
+import pandas as pd
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--magnification', type=str, default='20x')
+    parser.add_argument('--dest_weights', type=str)
+    args = parser.parse_args()
+    config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader)
+    dataset = DataSetWrapper(config['batch_size'], **config['dataset'])
+    
+    simclr = SimCLR(dataset, config, args)
+    simclr.train()
+
+
+if __name__ == "__main__":
+    main()

feature_extractor/simclr.py

@@ -0,0 +1,165 @@
+import torch
+from models.resnet_simclr import ResNetSimCLR
+from torch.utils.tensorboard import SummaryWriter
+import torch.nn.functional as F
+from loss.nt_xent import NTXentLoss
+import os
+import shutil
+import sys
+
+apex_support = False
+try:
+    sys.path.append('./apex')
+    from apex import amp
+
+    apex_support = True
+except:
+    print("Please install apex for mixed precision training from: https://github.com/NVIDIA/apex")
+    apex_support = False
+
+import numpy as np
+
+torch.manual_seed(0)
+
+
+def _save_config_file(model_checkpoints_folder):
+    if not os.path.exists(model_checkpoints_folder):
+        os.makedirs(model_checkpoints_folder)
+        shutil.copy('./config.yaml', os.path.join(model_checkpoints_folder, 'config.yaml'))
+
+
+class SimCLR(object):
+
+    def __init__(self, dataset, config, args=None):
+        self.config = config
+        self.device = self._get_device()
+        self.writer = SummaryWriter()
+        self.dataset = dataset
+        self.nt_xent_criterion = NTXentLoss(self.device, config['batch_size'], **config['loss'])
+        self.args = args
+    def _get_device(self):
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print("Running on:", device)
+        return device
+
+    def _step(self, model, xis, xjs, n_iter):
+
+        # get the representations and the projections
+        ris, zis = model(xis)  # [N,C]
+
+        # get the representations and the projections
+        rjs, zjs = model(xjs)  # [N,C]
+
+        # normalize projection feature vectors
+        zis = F.normalize(zis, dim=1)
+        zjs = F.normalize(zjs, dim=1)
+
+        loss = self.nt_xent_criterion(zis, zjs)
+        return loss
+
+    def train(self):
+
+        train_loader, valid_loader = self.dataset.get_data_loaders()
+
+        model = ResNetSimCLR(**self.config["model"])# .to(self.device)
+        if self.config['n_gpu'] > 1:
+            model = torch.nn.DataParallel(model, device_ids=eval(self.config['gpu_ids']))
+        model = self._load_pre_trained_weights(model)
+        model = model.to(self.device)
+            
+
+        optimizer = torch.optim.Adam(model.parameters(), 1e-5, weight_decay=eval(self.config['weight_decay']))
+
+#         scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0,
+#                                                                last_epoch=-1)
+        
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.config['epochs'], eta_min=0,
+                                                               last_epoch=-1)
+        
+
+        if apex_support and self.config['fp16_precision']:
+            model, optimizer = amp.initialize(model, optimizer,
+                                              opt_level='O2',
+                                              keep_batchnorm_fp32=True)
+
+        if self.args is None: 
+            model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints')
+        else:
+            model_checkpoints_folder = self.args.dest_weights#os.environ['FEATURE_EXTRACTOR_WEIGHT_PATH']
+            model_checkpoints_folder = os.path.dirname(model_checkpoints_folder)
+        # save config file
+        _save_config_file(model_checkpoints_folder)
+
+        n_iter = 0
+        valid_n_iter = 0
+        best_valid_loss = np.inf
+        
+        for epoch_counter in range(self.config['epochs']):
+            for (xis, xjs) in train_loader:
+                optimizer.zero_grad()
+                xis = xis.to(self.device)
+                xjs = xjs.to(self.device)
+
+                loss = self._step(model, xis, xjs, n_iter)
+
+                if n_iter % self.config['log_every_n_steps'] == 0:
+                    self.writer.add_scalar('train_loss', loss, global_step=n_iter)
+                    print("[%d/%d] step: %d train_loss: %.3f" % (epoch_counter, self.config['epochs'], n_iter, loss))
+
+                if apex_support and self.config['fp16_precision']:
+                    with amp.scale_loss(loss, optimizer) as scaled_loss:
+                        scaled_loss.backward()
+                else:
+                    loss.backward()
+
+                optimizer.step()
+                n_iter += 1
+
+            # validate the model if requested
+            if epoch_counter % self.config['eval_every_n_epochs'] == 0:
+                valid_loss = self._validate(model, valid_loader)
+                print("[%d/%d] val_loss: %.3f" % (epoch_counter, self.config['epochs'], valid_loss))
+                if valid_loss < best_valid_loss:
+                    # save the model weights
+                    best_valid_loss = valid_loss
+                    torch.save(model.state_dict(), os.path.join(model_checkpoints_folder, 'model.pth'))
+                    print('saved')
+
+                self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter)
+                valid_n_iter += 1
+
+            # warmup for the first 10 epochs
+            if epoch_counter >= 10:
+                scheduler.step()
+            self.writer.add_scalar('cosine_lr_decay', scheduler.get_lr()[0], global_step=n_iter)
+
+    def _load_pre_trained_weights(self, model):
+        try:
+            checkpoints_folder = os.path.join('./runs', self.config['fine_tune_from'], 'checkpoints')
+            state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth'))
+            model.load_state_dict(state_dict)
+            print("Loaded pre-trained model with success.")
+        except FileNotFoundError:
+            print("Pre-trained weights not found. Training from scratch.")
+
+        return model
+
+    def _validate(self, model, valid_loader):
+
+        # validation steps
+        with torch.no_grad():
+            model.eval()
+
+            valid_loss = 0.0
+            counter = 0
+
+            for (xis, xjs) in valid_loader:
+                xis = xis.to(self.device)
+                xjs = xjs.to(self.device)
+
+                loss = self._step(model, xis, xjs, counter)
+                valid_loss += loss.item()
+                counter += 1
+            valid_loss /= counter
+        model.train()
+        return valid_loss

feature_extractor/viewer.py

@@ -0,0 +1,227 @@
+#!/usr/bin/env python
+#
+# deepzoom_server - Example web application for serving whole-slide images
+#
+# Copyright (c) 2010-2015 Carnegie Mellon University
+#
+# This library is free software; you can redistribute it and/or modify it
+# under the terms of version 2.1 of the GNU Lesser General Public License
+# as published by the Free Software Foundation.
+#
+# This library is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+# License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this library; if not, write to the Free Software Foundation,
+# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+#
+
+from io import BytesIO
+from optparse import OptionParser
+import os
+import re
+from unicodedata import normalize
+
+from flask import Flask, abort, make_response, render_template, url_for
+
+if os.name == 'nt':
+    _dll_path = os.getenv('OPENSLIDE_PATH')
+    if _dll_path is not None:
+        if hasattr(os, 'add_dll_directory'):
+            # Python >= 3.8
+            with os.add_dll_directory(_dll_path):
+                import openslide
+        else:
+            # Python < 3.8
+            _orig_path = os.environ.get('PATH', '')
+            os.environ['PATH'] = _orig_path + ';' + _dll_path
+            import openslide
+
+            os.environ['PATH'] = _orig_path
+else:
+    import openslide
+
+from openslide import ImageSlide, open_slide
+from openslide.deepzoom import DeepZoomGenerator
+
+DEEPZOOM_SLIDE = None
+DEEPZOOM_FORMAT = 'jpeg'
+DEEPZOOM_TILE_SIZE = 254
+DEEPZOOM_OVERLAP = 1
+DEEPZOOM_LIMIT_BOUNDS = True
+DEEPZOOM_TILE_QUALITY = 75
+SLIDE_NAME = 'slide'
+
+app = Flask(__name__)
+app.config.from_object(__name__)
+app.config.from_envvar('DEEPZOOM_TILER_SETTINGS', silent=True)
+
+
+@app.before_first_request
+def load_slide():
+    slidefile = app.config['DEEPZOOM_SLIDE']
+    if slidefile is None:
+        raise ValueError('No slide file specified')
+    config_map = {
+        'DEEPZOOM_TILE_SIZE': 'tile_size',
+        'DEEPZOOM_OVERLAP': 'overlap',
+        'DEEPZOOM_LIMIT_BOUNDS': 'limit_bounds',
+    }
+    opts = {v: app.config[k] for k, v in config_map.items()}
+    slide = open_slide(slidefile)
+    app.slides = {SLIDE_NAME: DeepZoomGenerator(slide, **opts)}
+    app.associated_images = []
+    app.slide_properties = slide.properties
+    for name, image in slide.associated_images.items():
+        app.associated_images.append(name)
+        slug = slugify(name)
+        app.slides[slug] = DeepZoomGenerator(ImageSlide(image), **opts)
+    try:
+        mpp_x = slide.properties[openslide.PROPERTY_NAME_MPP_X]
+        mpp_y = slide.properties[openslide.PROPERTY_NAME_MPP_Y]
+        app.slide_mpp = (float(mpp_x) + float(mpp_y)) / 2
+    except (KeyError, ValueError):
+        app.slide_mpp = 0
+
+
+@app.route('/')
+def index():
+    slide_url = url_for('dzi', slug=SLIDE_NAME)
+    associated_urls = {
+        name: url_for('dzi', slug=slugify(name)) for name in app.associated_images
+    }
+    return render_template(
+        'slide-multipane.html',
+        slide_url=slide_url,
+        associated=associated_urls,
+        properties=app.slide_properties,
+        slide_mpp=app.slide_mpp,
+    )
+
+
+@app.route('/<slug>.dzi')
+def dzi(slug):
+    format = app.config['DEEPZOOM_FORMAT']
+    try:
+        resp = make_response(app.slides[slug].get_dzi(format))
+        resp.mimetype = 'application/xml'
+        return resp
+    except KeyError:
+        # Unknown slug
+        abort(404)
+
+
+@app.route('/<slug>_files/<int:level>/<int:col>_<int:row>.<format>')
+def tile(slug, level, col, row, format):
+    format = format.lower()
+    if format != 'jpeg' and format != 'png':
+        # Not supported by Deep Zoom
+        abort(404)
+    try:
+        tile = app.slides[slug].get_tile(level, (col, row))
+    except KeyError:
+        # Unknown slug
+        abort(404)
+    except ValueError:
+        # Invalid level or coordinates
+        abort(404)
+    buf = BytesIO()
+    tile.save(buf, format, quality=app.config['DEEPZOOM_TILE_QUALITY'])
+    resp = make_response(buf.getvalue())
+    resp.mimetype = 'image/%s' % format
+    return resp
+
+
+def slugify(text):
+    text = normalize('NFKD', text.lower()).encode('ascii', 'ignore').decode()
+    return re.sub('[^a-z0-9]+', '-', text)
+
+
+if __name__ == '__main__':
+    parser = OptionParser(usage='Usage: %prog [options] [slide]')
+    parser.add_option(
+        '-B',
+        '--ignore-bounds',
+        dest='DEEPZOOM_LIMIT_BOUNDS',
+        default=True,
+        action='store_false',
+        help='display entire scan area',
+    )
+    parser.add_option(
+        '-c', '--config', metavar='FILE', dest='config', help='config file'
+    )
+    parser.add_option(
+        '-d',
+        '--debug',
+        dest='DEBUG',
+        action='store_true',
+        help='run in debugging mode (insecure)',
+    )
+    parser.add_option(
+        '-e',
+        '--overlap',
+        metavar='PIXELS',
+        dest='DEEPZOOM_OVERLAP',
+        type='int',
+        help='overlap of adjacent tiles [1]',
+    )
+    parser.add_option(
+        '-f',
+        '--format',
+        metavar='{jpeg|png}',
+        dest='DEEPZOOM_FORMAT',
+        help='image format for tiles [jpeg]',
+    )
+    parser.add_option(
+        '-l',
+        '--listen',
+        metavar='ADDRESS',
+        dest='host',
+        default='127.0.0.1',
+        help='address to listen on [127.0.0.1]',
+    )
+    parser.add_option(
+        '-p',
+        '--port',
+        metavar='PORT',
+        dest='port',
+        type='int',
+        default=5000,
+        help='port to listen on [5000]',
+    )
+    parser.add_option(
+        '-Q',
+        '--quality',
+        metavar='QUALITY',
+        dest='DEEPZOOM_TILE_QUALITY',
+        type='int',
+        help='JPEG compression quality [75]',
+    )
+    parser.add_option(
+        '-s',
+        '--size',
+        metavar='PIXELS',
+        dest='DEEPZOOM_TILE_SIZE',
+        type='int',
+        help='tile size [254]',
+    )
+
+    (opts, args) = parser.parse_args()
+    # Load config file if specified
+    if opts.config is not None:
+        app.config.from_pyfile(opts.config)
+    # Overwrite only those settings specified on the command line
+    for k in dir(opts):
+        if not k.startswith('_') and getattr(opts, k) is None:
+            delattr(opts, k)
+    app.config.from_object(opts)
+    # Set slide file
+    try:
+        app.config['DEEPZOOM_SLIDE'] = args[0]
+    except IndexError:
+        if app.config['DEEPZOOM_SLIDE'] is None:
+            parser.error('No slide file specified')
+
+    app.run(host=opts.host, port=opts.port, threaded=True)

helper.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from __future__ import absolute_import, division, print_function
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from torchvision import transforms
+from utils.metrics import ConfusionMatrix
+from PIL import Image
+import os
+
+# torch.cuda.synchronize()
+# torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.deterministic = True
+
+def collate(batch):
+    image = [ b['image'] for b in batch ] # w, h
+    label = [ b['label'] for b in batch ]
+    id = [ b['id'] for b in batch ]
+    adj_s = [ b['adj_s'] for b in batch ]
+    return {'image': image, 'label': label, 'id': id, 'adj_s': adj_s}
+
+def preparefeatureLabel(batch_graph, batch_label, batch_adjs, device='cpu'):
+    batch_size = len(batch_graph)
+    labels = torch.LongTensor(batch_size)
+    max_node_num = 0
+
+    for i in range(batch_size):
+        labels[i] = batch_label[i]
+        max_node_num = max(max_node_num, batch_graph[i].shape[0])
+    
+    masks = torch.zeros(batch_size, max_node_num)
+    adjs =  torch.zeros(batch_size, max_node_num, max_node_num)
+    batch_node_feat = torch.zeros(batch_size, max_node_num, 512)
+
+    for i in range(batch_size):
+        cur_node_num =  batch_graph[i].shape[0]
+        #node attribute feature
+        tmp_node_fea = batch_graph[i]
+        batch_node_feat[i, 0:cur_node_num] = tmp_node_fea
+
+        #adjs
+        adjs[i, 0:cur_node_num, 0:cur_node_num] = batch_adjs[i]
+        
+        #masks
+        masks[i,0:cur_node_num] = 1  
+
+    node_feat = batch_node_feat.to(device)
+    labels = labels.to(device)
+    adjs = adjs.to(device)
+    masks = masks.to(device)
+
+    return node_feat, labels, adjs, masks
+
+class Trainer(object):
+    def __init__(self, n_class):
+        self.metrics = ConfusionMatrix(n_class)
+
+    def get_scores(self):
+        acc = self.metrics.get_scores()
+
+        return acc
+
+    def reset_metrics(self):
+        self.metrics.reset()
+    
+    def plot_cm(self):
+        self.metrics.plotcm()
+
+    def train(self, sample, model):
+        node_feat, labels, adjs, masks = preparefeatureLabel(sample['image'], sample['label'], sample['adj_s'])
+        pred,labels,loss = model.forward(node_feat, labels, adjs, masks)
+
+        return pred,labels,loss
+
+class Evaluator(object):
+    def __init__(self, n_class):
+        self.metrics = ConfusionMatrix(n_class)
+    
+    def get_scores(self):
+        acc = self.metrics.get_scores()
+
+        return acc
+
+    def reset_metrics(self):
+        self.metrics.reset()
+
+    def plot_cm(self):
+        self.metrics.plotcm()
+
+    def eval_test(self, sample, model, graphcam_flag=False):
+        node_feat, labels, adjs, masks = preparefeatureLabel(sample['image'], sample['label'], sample['adj_s'])
+        if not graphcam_flag:
+            with torch.no_grad():
+                pred,labels,loss = model.forward(node_feat, labels, adjs, masks)
+        else:
+            torch.set_grad_enabled(True)
+            pred,labels,loss= model.forward(node_feat, labels, adjs, masks, graphcam_flag=graphcam_flag)
+        return pred,labels,loss

main.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from __future__ import absolute_import, division, print_function
+
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from utils.dataset import GraphDataset
+from utils.lr_scheduler import LR_Scheduler
+from tensorboardX import SummaryWriter
+from helper import Trainer, Evaluator, collate
+from option import Options
+
+from models.GraphTransformer import Classifier
+from models.weight_init import weight_init
+import pickle
+args = Options().parse()
+
+label_map = pickle.load(open(os.path.join(args.dataset_metadata_path, 'label_map.pkl'), 'rb'))
+
+n_class = len(label_map)
+
+torch.cuda.synchronize()
+torch.backends.cudnn.deterministic = True
+
+data_path = args.data_path
+model_path = args.model_path
+if not os.path.isdir(model_path): os.mkdir(model_path)
+log_path = args.log_path
+if not os.path.isdir(log_path): os.mkdir(log_path)
+task_name = args.task_name
+
+print(task_name)
+###################################
+train = args.train
+test = args.test
+graphcam = args.graphcam
+print("train:", train, "test:", test, "graphcam:", graphcam)
+
+##### Load datasets
+print("preparing datasets and dataloaders......")
+batch_size = args.batch_size
+
+if train:
+    ids_train = open(args.train_set).readlines()
+    dataset_train = GraphDataset(os.path.join(data_path, ""), ids_train, args.dataset_metadata_path)
+    dataloader_train = torch.utils.data.DataLoader(dataset=dataset_train, batch_size=batch_size, num_workers=10, collate_fn=collate, shuffle=True, pin_memory=True, drop_last=True)
+    total_train_num = len(dataloader_train) * batch_size
+
+ids_val = open(args.val_set).readlines()
+dataset_val = GraphDataset(os.path.join(data_path, ""), ids_val, args.dataset_metadata_path)
+dataloader_val = torch.utils.data.DataLoader(dataset=dataset_val, batch_size=batch_size, num_workers=10, collate_fn=collate, shuffle=False, pin_memory=True)
+total_val_num = len(dataloader_val) * batch_size
+    
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+##### creating models #############
+print("creating models......")
+
+num_epochs = args.num_epochs
+learning_rate = args.lr
+
+model = Classifier(n_class)
+model = nn.DataParallel(model)
+if args.resume:
+    print('load model{}'.format(args.resume))
+    model.load_state_dict(torch.load(args.resume))
+
+if torch.cuda.is_available():
+    model = model.cuda()
+#model.apply(weight_init)
+
+optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate, weight_decay = 5e-4)       # best:5e-4, 4e-3
+scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[20,100], gamma=0.1) # gamma=0.3  # 30,90,130 # 20,90,130 -> 150
+
+##################################
+
+criterion = nn.CrossEntropyLoss()
+
+if not test:
+    writer = SummaryWriter(log_dir=log_path + task_name)
+    f_log = open(log_path + task_name + ".log", 'w')
+
+trainer = Trainer(n_class)
+evaluator = Evaluator(n_class)
+
+best_pred = 0.0
+for epoch in range(num_epochs):
+    # optimizer.zero_grad()
+    model.train()
+    train_loss = 0.
+    total = 0.
+
+    current_lr = optimizer.param_groups[0]['lr']
+    print('\n=>Epoches %i, learning rate = %.7f, previous best = %.4f' % (epoch+1, current_lr, best_pred))
+
+    if train:
+        for i_batch, sample_batched in enumerate(dataloader_train):
+            scheduler.step(epoch)
+
+            preds,labels,loss = trainer.train(sample_batched, model)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss
+            total += len(labels)
+
+            trainer.metrics.update(labels, preds)
+            if (i_batch + 1) % args.log_interval_local == 0:
+                print("[%d/%d] train loss: %.3f; agg acc: %.3f" % (total, total_train_num, train_loss / total, trainer.get_scores()))
+                trainer.plot_cm()
+
+    if not test: 
+        print("[%d/%d] train loss: %.3f; agg acc: %.3f" % (total_train_num, total_train_num, train_loss / total, trainer.get_scores()))
+        trainer.plot_cm()
+
+
+    if epoch % 1 == 0:
+        with torch.no_grad():
+            model.eval()
+            print("evaluating...")
+
+            total = 0.
+            batch_idx = 0
+
+            for i_batch, sample_batched in enumerate(dataloader_val):
+                preds, labels, _ = evaluator.eval_test(sample_batched, model, graphcam)
+                
+                total += len(labels)
+
+                evaluator.metrics.update(labels, preds)
+
+                if (i_batch + 1) % args.log_interval_local == 0:
+                    print('[%d/%d] val agg acc: %.3f' % (total, total_val_num, evaluator.get_scores()))
+                    evaluator.plot_cm()
+
+            print('[%d/%d] val agg acc: %.3f' % (total_val_num, total_val_num, evaluator.get_scores()))
+            evaluator.plot_cm()
+
+            # torch.cuda.empty_cache()
+
+            val_acc = evaluator.get_scores()
+            if val_acc > best_pred: 
+                best_pred = val_acc
+                if not test:
+                    print("saving model...")
+                    torch.save(model.state_dict(), model_path + task_name + ".pth")
+
+            log = ""
+            log = log + 'epoch [{}/{}] ------ acc: train = {:.4f}, val = {:.4f}'.format(epoch+1, num_epochs, trainer.get_scores(), evaluator.get_scores()) + "\n"
+
+            log += "================================\n"
+            print(log)
+            if test: break
+
+            f_log.write(log)
+            f_log.flush()
+
+            writer.add_scalars('accuracy', {'train acc': trainer.get_scores(), 'val acc': evaluator.get_scores()}, epoch+1)
+
+    trainer.reset_metrics()
+    evaluator.reset_metrics()
+
+if not test: f_log.close()

metadata/label_map.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce5be416a8667c9379502eaf8407e6d07bbae03749085190be630bd3b026eb52
+size 34

models/.gitkeep

@@ -0,0 +1 @@
+

models/GraphTransformer.py

@@ -0,0 +1,123 @@
+import sys
+import os
+import torch
+import random
+import numpy as np
+
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+from .ViT import *
+from .gcn import GCNBlock
+
+from torch_geometric.nn import GCNConv, DenseGraphConv, dense_mincut_pool
+from torch.nn import Linear
+class Classifier(nn.Module):
+    def __init__(self, n_class):
+        super(Classifier, self).__init__()
+
+        self.n_class = n_class
+        self.embed_dim = 64
+        self.num_layers = 3
+        self.node_cluster_num = 100
+
+        self.transformer = VisionTransformer(num_classes=n_class, embed_dim=self.embed_dim)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+        self.criterion = nn.CrossEntropyLoss()
+
+        self.bn = 1
+        self.add_self = 1
+        self.normalize_embedding = 1
+        self.conv1 = GCNBlock(512,self.embed_dim,self.bn,self.add_self,self.normalize_embedding,0.,0)       # 64->128
+        self.pool1 = Linear(self.embed_dim, self.node_cluster_num)                                          # 100-> 20
+
+
+    def forward(self,node_feat,labels,adj,mask,is_print=False, graphcam_flag=False, to_file=True):
+        # node_feat, labels = self.PrepareFeatureLabel(batch_graph)
+        cls_loss=node_feat.new_zeros(self.num_layers)
+        rank_loss=node_feat.new_zeros(self.num_layers-1)
+        X=node_feat
+        p_t=[]
+        pred_logits=0
+        visualize_tools=[]
+        if labels is not None:
+            visualize_tools1=[labels.cpu()]
+        embeds=0
+        concats=[]
+        
+        layer_acc=[]
+                
+        X=mask.unsqueeze(2)*X
+        X = self.conv1(X, adj, mask)
+        s = self.pool1(X)
+
+
+        graphcam_tensors = {}
+
+        if graphcam_flag:
+            s_matrix = torch.argmax(s[0], dim=1)
+            if to_file:
+                from os import path
+                os.makedirs('graphcam', exist_ok=True)
+                torch.save(s_matrix, 'graphcam/s_matrix.pt')
+                torch.save(s[0], 'graphcam/s_matrix_ori.pt')
+                
+                if path.exists('graphcam/att_1.pt'):
+                    os.remove('graphcam/att_1.pt')
+                    os.remove('graphcam/att_2.pt')
+                    os.remove('graphcam/att_3.pt')
+
+            if not to_file:
+                graphcam_tensors['s_matrix'] = s_matrix
+                graphcam_tensors['s_matrix_ori'] = s[0]
+
+    
+        X, adj, mc1, o1 = dense_mincut_pool(X, adj, s, mask)
+        b, _, _ = X.shape
+        cls_token = self.cls_token.repeat(b, 1, 1)
+        X = torch.cat([cls_token, X], dim=1)
+
+        out = self.transformer(X)
+
+        loss = None
+        if labels is not None:
+            # loss
+            loss = self.criterion(out, labels)
+            loss = loss + mc1 + o1
+        # pred
+        pred = out.data.max(1)[1]
+
+        if graphcam_flag:
+            #print('GraphCAM enabled')
+            #print(out.shape)
+            p = F.softmax(out)
+            #print(p.shape)
+            if to_file:
+                torch.save(p, 'graphcam/prob.pt')
+            if not to_file:
+                graphcam_tensors['prob'] = p          
+            index = np.argmax(out.cpu().data.numpy(), axis=-1)
+
+            for index_ in range(self.n_class):
+                one_hot = np.zeros((1, out.size()[-1]), dtype=np.float32)
+                one_hot[0, index_] = out[0][index_]
+                one_hot_vector = one_hot
+                one_hot = torch.from_numpy(one_hot).requires_grad_(True)
+                one_hot = torch.sum(one_hot.to( 'cuda' if torch.cuda.is_available() else 'cpu') * out)       #!!!!!!!!!!!!!!!!!!!!out-->p
+                self.transformer.zero_grad()
+                one_hot.backward(retain_graph=True)
+
+                kwargs = {"alpha": 1}
+                cam = self.transformer.relprop(torch.tensor(one_hot_vector).to(X.device), method="transformer_attribution", is_ablation=False, 
+                                            start_layer=0, **kwargs)
+                if to_file:
+                    torch.save(cam, 'graphcam/cam_{}.pt'.format(index_))
+                if not to_file:
+                    graphcam_tensors[f'cam_{index_}'] = cam                
+        
+        if not to_file:
+            return pred,labels,loss, graphcam_tensors
+        return pred,labels,loss

models/ViT.py

@@ -0,0 +1,415 @@
+""" Vision Transformer (ViT) in PyTorch
+"""
+import torch
+import torch.nn as nn
+from einops import rearrange
+from .layers import *
+import math
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'first_conv': 'patch_embed.proj', 'classifier': 'head',
+        **kwargs
+    }
+
+
+default_cfgs = {
+    # patch models
+    'vit_small_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/vit_small_p16_224-15ec54c9.pth',
+    ),
+    'vit_base_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
+    ),
+    'vit_large_patch16_224': _cfg(
+        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+}
+
+def compute_rollout_attention(all_layer_matrices, start_layer=0):
+    # adding residual consideration
+    num_tokens = all_layer_matrices[0].shape[1]
+    batch_size = all_layer_matrices[0].shape[0]
+    eye = torch.eye(num_tokens).expand(batch_size, num_tokens, num_tokens).to(all_layer_matrices[0].device)
+    all_layer_matrices = [all_layer_matrices[i] + eye for i in range(len(all_layer_matrices))]
+    # all_layer_matrices = [all_layer_matrices[i] / all_layer_matrices[i].sum(dim=-1, keepdim=True)
+    #                       for i in range(len(all_layer_matrices))]
+    joint_attention = all_layer_matrices[start_layer]
+    for i in range(start_layer+1, len(all_layer_matrices)):
+        joint_attention = all_layer_matrices[i].bmm(joint_attention)
+    return joint_attention
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = Linear(in_features, hidden_features)
+        self.act = GELU()
+        self.fc2 = Linear(hidden_features, out_features)
+        self.drop = Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+    def relprop(self, cam, **kwargs):
+        cam = self.drop.relprop(cam, **kwargs)
+        cam = self.fc2.relprop(cam, **kwargs)
+        cam = self.act.relprop(cam, **kwargs)
+        cam = self.fc1.relprop(cam, **kwargs)
+        return cam
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False,attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = head_dim ** -0.5
+
+        # A = Q*K^T
+        self.matmul1 = einsum('bhid,bhjd->bhij')
+        # attn = A*V
+        self.matmul2 = einsum('bhij,bhjd->bhid')
+
+        self.qkv = Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = Dropout(attn_drop)
+        self.proj = Linear(dim, dim)
+        self.proj_drop = Dropout(proj_drop)
+        self.softmax = Softmax(dim=-1)
+
+        self.attn_cam = None
+        self.attn = None
+        self.v = None
+        self.v_cam = None
+        self.attn_gradients = None
+
+    def get_attn(self):
+        return self.attn
+
+    def save_attn(self, attn):
+        self.attn = attn
+
+    def save_attn_cam(self, cam):
+        self.attn_cam = cam
+
+    def get_attn_cam(self):
+        return self.attn_cam
+
+    def get_v(self):
+        return self.v
+
+    def save_v(self, v):
+        self.v = v
+
+    def save_v_cam(self, cam):
+        self.v_cam = cam
+
+    def get_v_cam(self):
+        return self.v_cam
+
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+
+    def get_attn_gradients(self):
+        return self.attn_gradients
+
+    def forward(self, x):
+        b, n, _, h = *x.shape, self.num_heads
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=h)
+
+        self.save_v(v)
+
+        dots = self.matmul1([q, k]) * self.scale
+
+        attn = self.softmax(dots)
+        attn = self.attn_drop(attn)
+
+        # Get attention
+        if False:
+            from os import path
+            if not path.exists('att_1.pt'):
+                torch.save(attn, 'att_1.pt')
+            elif not path.exists('att_2.pt'):
+                torch.save(attn, 'att_2.pt')
+            else:
+                torch.save(attn, 'att_3.pt')
+
+        #comment in training
+        if x.requires_grad:
+            self.save_attn(attn)
+            attn.register_hook(self.save_attn_gradients)
+
+        out = self.matmul2([attn, v])
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        out = self.proj(out)
+        out = self.proj_drop(out)
+        return out
+
+    def relprop(self, cam, **kwargs):
+        cam = self.proj_drop.relprop(cam, **kwargs)
+        cam = self.proj.relprop(cam, **kwargs)
+        cam = rearrange(cam, 'b n (h d) -> b h n d', h=self.num_heads)
+
+        # attn = A*V
+        (cam1, cam_v)= self.matmul2.relprop(cam, **kwargs)
+        cam1 /= 2
+        cam_v /= 2
+
+        self.save_v_cam(cam_v)
+        self.save_attn_cam(cam1)
+
+        cam1 = self.attn_drop.relprop(cam1, **kwargs)
+        cam1 = self.softmax.relprop(cam1, **kwargs)
+
+        # A = Q*K^T
+        (cam_q, cam_k) = self.matmul1.relprop(cam1, **kwargs)
+        cam_q /= 2
+        cam_k /= 2
+
+        cam_qkv = rearrange([cam_q, cam_k, cam_v], 'qkv b h n d -> b n (qkv h d)', qkv=3, h=self.num_heads)
+
+        return self.qkv.relprop(cam_qkv, **kwargs)
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.):
+        super().__init__()
+        self.norm1 = LayerNorm(dim, eps=1e-6)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        self.norm2 = LayerNorm(dim, eps=1e-6)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)
+
+        self.add1 = Add()
+        self.add2 = Add()
+        self.clone1 = Clone()
+        self.clone2 = Clone()
+
+    def forward(self, x):
+        x1, x2 = self.clone1(x, 2)
+        x = self.add1([x1, self.attn(self.norm1(x2))])
+        x1, x2 = self.clone2(x, 2)
+        x = self.add2([x1, self.mlp(self.norm2(x2))])
+        return x
+
+    def relprop(self, cam, **kwargs):
+        (cam1, cam2) = self.add2.relprop(cam, **kwargs)
+        cam2 = self.mlp.relprop(cam2, **kwargs)
+        cam2 = self.norm2.relprop(cam2, **kwargs)
+        cam = self.clone2.relprop((cam1, cam2), **kwargs)
+
+        (cam1, cam2) = self.add1.relprop(cam, **kwargs)
+        cam2 = self.attn.relprop(cam2, **kwargs)
+        cam2 = self.norm1.relprop(cam2, **kwargs)
+        cam = self.clone1.relprop((cam1, cam2), **kwargs)
+        return cam
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, num_classes=2, embed_dim=64, depth=3,
+                 num_heads=8, mlp_ratio=2., qkv_bias=False, mlp_head=False, drop_rate=0., attn_drop_rate=0.):
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                drop=drop_rate, attn_drop=attn_drop_rate)
+            for i in range(depth)])
+
+        self.norm = LayerNorm(embed_dim)
+        if mlp_head:
+            # paper diagram suggests 'MLP head', but results in 4M extra parameters vs paper
+            self.head = Mlp(embed_dim, int(embed_dim * mlp_ratio), num_classes)
+        else:
+            # with a single Linear layer as head, the param count within rounding of paper
+            self.head = Linear(embed_dim, num_classes)
+
+        #self.apply(self._init_weights)
+
+        self.pool = IndexSelect()
+        self.add = Add()
+
+        self.inp_grad = None
+
+    def save_inp_grad(self,grad):
+        self.inp_grad = grad
+
+    def get_inp_grad(self):
+        return self.inp_grad
+
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @property
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def forward(self, x):
+        if x.requires_grad:
+            x.register_hook(self.save_inp_grad)     #comment it in train
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = self.norm(x)
+        x = self.pool(x, dim=1, indices=torch.tensor(0, device=x.device))
+        x = x.squeeze(1)
+        x = self.head(x)
+        return x
+
+    def relprop(self, cam=None,method="transformer_attribution", is_ablation=False, start_layer=0, **kwargs):
+        # print(kwargs)
+        # print("conservation 1", cam.sum())
+        cam = self.head.relprop(cam, **kwargs)
+        cam = cam.unsqueeze(1)
+        cam = self.pool.relprop(cam, **kwargs)
+        cam = self.norm.relprop(cam, **kwargs)
+        for blk in reversed(self.blocks):
+            cam = blk.relprop(cam, **kwargs)
+
+        # print("conservation 2", cam.sum())
+        # print("min", cam.min())
+
+        if method == "full":
+            (cam, _) = self.add.relprop(cam, **kwargs)
+            cam = cam[:, 1:]
+            cam = self.patch_embed.relprop(cam, **kwargs)
+            # sum on channels
+            cam = cam.sum(dim=1)
+            return cam
+
+        elif method == "rollout":
+            # cam rollout
+            attn_cams = []
+            for blk in self.blocks:
+                attn_heads = blk.attn.get_attn_cam().clamp(min=0)
+                avg_heads = (attn_heads.sum(dim=1) / attn_heads.shape[1]).detach()
+                attn_cams.append(avg_heads)
+            cam = compute_rollout_attention(attn_cams, start_layer=start_layer)
+            cam = cam[:, 0, 1:]
+            return cam
+        
+        # our method, method name grad is legacy
+        elif method == "transformer_attribution" or method == "grad":
+            cams = []
+            for blk in self.blocks:
+                grad = blk.attn.get_attn_gradients()
+                cam = blk.attn.get_attn_cam()
+                cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
+                grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])
+                cam = grad * cam
+                cam = cam.clamp(min=0).mean(dim=0)
+                cams.append(cam.unsqueeze(0))
+            rollout = compute_rollout_attention(cams, start_layer=start_layer)
+            cam = rollout[:, 0, 1:]
+            return cam
+            
+        elif method == "last_layer":
+            cam = self.blocks[-1].attn.get_attn_cam()
+            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
+            if is_ablation:
+                grad = self.blocks[-1].attn.get_attn_gradients()
+                grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])
+                cam = grad * cam
+            cam = cam.clamp(min=0).mean(dim=0)
+            cam = cam[0, 1:]
+            return cam
+
+        elif method == "last_layer_attn":
+            cam = self.blocks[-1].attn.get_attn()
+            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
+            cam = cam.clamp(min=0).mean(dim=0)
+            cam = cam[0, 1:]
+            return cam
+
+        elif method == "second_layer":
+            cam = self.blocks[1].attn.get_attn_cam()
+            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
+            if is_ablation:
+                grad = self.blocks[1].attn.get_attn_gradients()
+                grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])
+                cam = grad * cam
+            cam = cam.clamp(min=0).mean(dim=0)
+            cam = cam[0, 1:]
+            return cam

models/gcn.py

@@ -0,0 +1,420 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import torch.nn.functional as F
+import math
+
+import numpy as np
+
+torch.set_printoptions(precision=2,threshold=float('inf'))
+
+class AGCNBlock(nn.Module):
+    def __init__(self,input_dim,hidden_dim,gcn_layer=2,dropout=0.0,relu=0):
+        super(AGCNBlock,self).__init__()
+        if dropout > 0.001:
+            self.dropout_layer = nn.Dropout(p=dropout)
+        self.sort = 'sort'
+        self.model='agcn'
+        self.gcns=nn.ModuleList()
+        self.bn = 0
+        self.add_self = 1
+        self.normalize_embedding = 1
+        self.gcns.append(GCNBlock(input_dim,hidden_dim,self.bn,self.add_self,self.normalize_embedding,dropout,relu))
+        self.pool = 'mean'
+        self.tau = 1.
+        self.lamda = 1.
+
+        for i in range(gcn_layer-1):
+            if i==gcn_layer-2 and (not 1):
+                self.gcns.append(GCNBlock(hidden_dim,hidden_dim,self.bn,self.add_self,self.normalize_embedding,dropout,0))
+            else:
+                self.gcns.append(GCNBlock(hidden_dim,hidden_dim,self.bn,self.add_self,self.normalize_embedding,dropout,relu))
+            
+        if self.model=='diffpool':
+            self.pool_gcns=nn.ModuleList()
+            tmp=input_dim
+            self.diffpool_k=200
+            for i in range(3):
+                self.pool_gcns.append(GCNBlock(tmp,200,0,0,0,dropout,relu))
+                tmp=200
+
+        self.w_a=nn.Parameter(torch.zeros(1,hidden_dim,1))
+        self.w_b=nn.Parameter(torch.zeros(1,hidden_dim,1))
+        torch.nn.init.normal_(self.w_a)
+        torch.nn.init.uniform_(self.w_b,-1,1)
+
+        self.pass_dim=hidden_dim
+
+        if self.pool=='mean':
+            self.pool=self.mean_pool
+        elif self.pool=='max':
+            self.pool=self.max_pool
+        elif self.pool=='sum':
+            self.pool=self.sum_pool
+
+        self.softmax='global'
+        if self.softmax=='gcn':
+            self.att_gcn=GCNBlock(2,1,0,0,dropout,relu)
+        self.khop=1
+        self.adj_norm='none'
+
+        self.filt_percent=0.25       #default 0.5
+        self.eps=1e-10
+
+        self.tau_config=1
+        if 1==-1.:
+            self.tau=nn.Parameter(torch.tensor(1),requires_grad=False)
+        elif 1==-2.:
+            self.tau_fc=nn.Linear(hidden_dim,1)
+            torch.nn.init.constant_(self.tau_fc.bias,1)
+            torch.nn.init.xavier_normal_(self.tau_fc.weight.t())
+        else:
+            self.tau=nn.Parameter(torch.tensor(self.tau))
+        self.lamda1=nn.Parameter(torch.tensor(self.lamda))
+        self.lamda2=nn.Parameter(torch.tensor(self.lamda))
+
+        self.att_norm=0
+        
+        self.dnorm=0
+        self.dnorm_coe=1
+
+        self.att_out=0
+        self.single_att=0
+
+
+    def forward(self,X,adj,mask,is_print=False):
+        '''
+    input:
+        X:  node input features , [batch,node_num,input_dim],dtype=float
+        adj: adj matrix, [batch,node_num,node_num], dtype=float
+        mask: mask for nodes, [batch,node_num]
+    outputs:
+        out:unormalized classification prob, [batch,hidden_dim]
+        H: batch of node hidden features, [batch,node_num,pass_dim]
+        new_adj: pooled new adj matrix, [batch, k_max, k_max]
+        new_mask: [batch, k_max]
+        '''
+        hidden=X
+        #adj = adj.float()
+        # print('input size:')
+        # print(hidden.shape)
+        
+        is_print1=is_print2=is_print
+        if adj.shape[-1]>100:
+            is_print1=False
+
+        for gcn in self.gcns:
+            hidden=gcn(hidden,adj,mask)
+        #     print('gcn:')
+        #     print(hidden.shape)
+        # print('mask:')
+        # print(mask.unsqueeze(2).shape)
+        # print(mask.sum(dim=1))
+
+        hidden=mask.unsqueeze(2)*hidden
+        # print(hidden[0][0])
+        # print(hidden[0][-1])
+
+        if self.model=='unet':
+            att=torch.matmul(hidden,self.w_a).squeeze()
+            att=att/torch.sqrt((self.w_a.squeeze(2)**2).sum(dim=1,keepdim=True))
+        elif self.model=='agcn':
+            if self.softmax=='global' or self.softmax=='mix':
+                if False:
+                    dgree_w = torch.sum(adj, dim=2) / torch.sum(adj, dim=2).max(1, keepdim=True)[0]
+                    att_a=torch.matmul(hidden,self.w_a).squeeze()*dgree_w+(mask-1)*1e10
+                else:
+                    att_a=torch.matmul(hidden,self.w_a).squeeze()+(mask-1)*1e10
+                    # print(att_a[0][:10])
+                    # print(att_a[0][-10:-1])
+                att_a_1=att_a=torch.nn.functional.softmax(att_a,dim=1)
+                # print(att_a[0][:10])
+                # print(att_a[0][-10:-1])
+
+                if self.dnorm:
+                    scale=mask.sum(dim=1,keepdim=True)/self.dnorm_coe
+                    att_a=scale*att_a
+            if self.softmax=='neibor' or self.softmax=='mix':
+                att_b=torch.matmul(hidden,self.w_b).squeeze()+(mask-1)*1e10
+                att_b_max,_=att_b.max(dim=1,keepdim=True)
+                if self.tau_config!=-2:
+                    att_b=torch.exp((att_b-att_b_max)*torch.abs(self.tau))
+                else:
+                    att_b=torch.exp((att_b-att_b_max)*torch.abs(self.tau_fc(self.pool(hidden,mask))))
+                denom=att_b.unsqueeze(2)
+                for _ in range(self.khop):
+                    denom=torch.matmul(adj,denom)
+                denom=denom.squeeze()+self.eps
+                att_b=(att_b*torch.diagonal(adj,0,1,2))/denom
+                if self.dnorm:
+                    if self.adj_norm=='diag':
+                        diag_scale=mask/(torch.diagonal(adj,0,1,2)+self.eps)
+                    elif self.adj_norm=='none':
+                        diag_scale=adj.sum(dim=1)
+                    att_b=att_b*diag_scale
+                att_b=att_b*mask
+                        
+            if self.softmax=='global':
+                att=att_a
+            elif self.softmax=='neibor' or self.softmax=='hardnei':
+                att=att_b
+            elif self.softmax=='mix':
+                att=att_a*torch.abs(self.lamda1)+att_b*torch.abs(self.lamda2)
+        # print('att:')
+        # print(att.shape)
+        Z=hidden
+
+        if self.model=='unet':
+            Z=torch.tanh(att.unsqueeze(2))*Z
+        elif self.model=='agcn':
+            if self.single_att:
+                Z=Z
+            else:
+                Z=att.unsqueeze(2)*Z
+        # print('Z shape')
+        # print(Z.shape)
+        k_max=int(math.ceil(self.filt_percent*adj.shape[-1]))
+        # print('k_max')
+        # print(k_max)
+        if self.model=='diffpool':
+            k_max=min(k_max,self.diffpool_k)
+
+        k_list=[int(math.ceil(self.filt_percent*x)) for x in mask.sum(dim=1).tolist()]
+        # print('k_list')
+        # print(k_list)
+        if self.model!='diffpool': 
+            if self.sort=='sample':
+                att_samp = att * mask
+                att_samp = (att_samp/att_samp.sum(1)).detach().cpu().numpy()
+                top_index = ()
+                for i in range(att.size(0)):
+                    top_index = (torch.LongTensor(np.random.choice(att_samp.size(1), k_max, att_samp[i])) ,)
+                top_index = torch.stack(top_index,1)
+            elif self.sort=='random_sample':
+                top_index = torch.LongTensor(att.size(0), k_max)*0
+                for i in range(att.size(0)):
+                    top_index[i,0:k_list[i]] = torch.randperm(int(mask[i].sum().item()))[0:k_list[i]] 
+            else: #sort
+                _,top_index=torch.topk(att,k_max,dim=1)
+        # print('top_index')
+        # print(top_index)
+        # print(len(top_index[0]))
+        new_mask=X.new_zeros(X.shape[0],k_max)
+        # print('new_mask')
+        # print(new_mask.shape)
+        visualize_tools=None 
+        if self.model=='unet':
+            for i,k in enumerate(k_list):
+                for j in range(int(k),k_max):
+                    top_index[i][j]=adj.shape[-1]-1
+                    new_mask[i][j]=-1.
+            new_mask=new_mask+1
+            top_index,_=torch.sort(top_index,dim=1)
+            assign_m=X.new_zeros(X.shape[0],k_max,adj.shape[-1])
+            for i,x in enumerate(top_index):
+                assign_m[i]=torch.index_select(adj[i],0,x)
+            new_adj=X.new_zeros(X.shape[0],k_max,k_max)
+            H=Z.new_zeros(Z.shape[0],k_max,Z.shape[-1])
+            for i,x in enumerate(top_index):
+                new_adj[i]=torch.index_select(assign_m[i],1,x)
+                H[i]=torch.index_select(Z[i],0,x)
+
+        elif self.model=='agcn':
+            assign_m=X.new_zeros(X.shape[0],k_max,adj.shape[-1])
+            # print('assign_m.shape')
+            # print(assign_m.shape)
+            for i,k in enumerate(k_list):
+                #print('top_index[i][j]')
+                for j in range(int(k)):  
+                    #print(str(top_index[i][j].item())+' ', end='')
+                    assign_m[i][j]=adj[i][top_index[i][j]]
+                    #print(assign_m[i][j])
+                    new_mask[i][j]=1.
+
+            assign_m=assign_m/(assign_m.sum(dim=1,keepdim=True)+self.eps)
+            H=torch.matmul(assign_m,Z)
+            # print('H')
+            # print(H.shape)
+            new_adj=torch.matmul(torch.matmul(assign_m,adj),torch.transpose(assign_m,1,2))
+            # print(torch.matmul(assign_m,adj).shape)
+            # print('new_adj:')
+            # print(new_adj.shape)
+            
+        elif self.model=='diffpool':
+            hidden1=X
+            for gcn in self.pool_gcns:
+                hidden1=gcn(hidden1,adj,mask)
+            assign_m=X.new_ones(X.shape[0],X.shape[1],k_max)*(-100000000.)
+            for i,x in enumerate(hidden1):
+                k=min(k_list[i],k_max)
+                assign_m[i,:,0:k]=hidden1[i,:,0:k]
+                for j in range(int(k)):
+                    new_mask[i][j]=1.
+
+            assign_m=torch.nn.functional.softmax(assign_m,dim=2)*mask.unsqueeze(2)
+            assign_m_t=torch.transpose(assign_m,1,2)
+            new_adj=torch.matmul(torch.matmul(assign_m_t,adj),assign_m)
+            H=torch.matmul(assign_m_t,Z)
+        # print('pool')    
+        if self.att_out and self.model=='agcn':
+            if self.softmax=='global':
+                out=self.pool(att_a_1.unsqueeze(2)*hidden,mask)
+            elif self.softmax=='neibor':
+                att_b_sum=att_b.sum(dim=1,keepdim=True)
+                out=self.pool((att_b/(att_b_sum+self.eps)).unsqueeze(2)*hidden,mask)
+        else:
+            # print('hidden.shape')
+            # print(hidden.shape)
+            out=self.pool(hidden,mask)
+            # print('out shape')
+            # print(out.shape)
+           
+        if self.adj_norm=='tanh' or self.adj_norm=='mix':
+            new_adj=torch.tanh(new_adj)
+        elif self.adj_norm=='diag' or self.adj_norm=='mix':
+            diag_elem=torch.pow(new_adj.sum(dim=2)+self.eps,-0.5)
+            diag=new_adj.new_zeros(new_adj.shape)
+            for i,x in enumerate(diag_elem):
+                diag[i]=torch.diagflat(x)
+            new_adj=torch.matmul(torch.matmul(diag,new_adj),diag)
+
+        visualize_tools=[]
+        '''
+        if (not self.training) and is_print1:
+            print('**********************************')
+            print('node_feat:',X.type(),X.shape)
+            print(X)
+            if self.model!='diffpool':
+                print('**********************************')
+                print('att:',att.type(),att.shape)
+                print(att)
+                print('**********************************')
+                print('top_index:',top_index.type(),top_index.shape)
+                print(top_index)
+            print('**********************************')
+            print('adj:',adj.type(),adj.shape)
+            print(adj)
+            print('**********************************')
+            print('assign_m:',assign_m.type(),assign_m.shape)
+            print(assign_m)
+            print('**********************************')
+            print('new_adj:',new_adj.type(),new_adj.shape)
+            print(new_adj)
+            print('**********************************')
+            print('new_mask:',new_mask.type(),new_mask.shape)
+            print(new_mask)
+        '''
+        #visualization
+        from os import path
+        if not path.exists('att_1.pt'):
+            torch.save(att[0], 'att_1.pt')
+            torch.save(top_index[0], 'att_ind1.pt')
+        elif not path.exists('att_2.pt'):
+            torch.save(att[0], 'att_2.pt')
+            torch.save(top_index[0], 'att_ind2.pt')
+        else:
+            torch.save(att[0], 'att_3.pt')
+            torch.save(top_index[0], 'att_ind3.pt')
+
+        if (not self.training) and is_print2:
+            if self.model!='diffpool':
+                visualize_tools.append(att[0])
+                visualize_tools.append(top_index[0])
+            visualize_tools.append(new_adj[0])
+            visualize_tools.append(new_mask.sum())
+        # print('**********************************')
+        return out,H,new_adj,new_mask,visualize_tools
+    
+    def mean_pool(self,x,mask):
+        return x.sum(dim=1)/(self.eps+mask.sum(dim=1,keepdim=True))
+    
+    def sum_pool(self,x,mask):
+        return x.sum(dim=1)
+
+    @staticmethod
+    def max_pool(x,mask):
+        #output: [batch,x.shape[2]]
+        m=(mask-1)*1e10
+        r,_=(x+m.unsqueeze(2)).max(dim=1)
+        return r
+# GCN basic operation
+class GCNBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, bn=0,add_self=0, normalize_embedding=0,
+            dropout=0.0,relu=0, bias=True):
+        super(GCNBlock,self).__init__()
+        self.add_self = add_self
+        self.dropout = dropout
+        self.relu=relu
+        self.bn=bn
+        if dropout > 0.001:
+            self.dropout_layer = nn.Dropout(p=dropout)
+        if self.bn:
+            self.bn_layer = torch.nn.BatchNorm1d(output_dim)
+
+        self.normalize_embedding = normalize_embedding
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+
+        self.weight = nn.Parameter(torch.FloatTensor(input_dim, output_dim).to( 'cuda' if torch.cuda.is_available() else 'cpu') )
+        torch.nn.init.xavier_normal_(self.weight)
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(output_dim).to( 'cuda' if torch.cuda.is_available() else 'cpu') )
+        else:
+            self.bias = None
+
+    def forward(self, x, adj, mask):
+        y = torch.matmul(adj, x)
+        if self.add_self:
+            y += x
+        y = torch.matmul(y,self.weight)
+        if self.bias is not None:
+            y = y + self.bias
+        if self.normalize_embedding:
+            y = F.normalize(y, p=2, dim=2)
+        if self.bn:
+            index=mask.sum(dim=1).long().tolist()
+            bn_tensor_bf=mask.new_zeros((sum(index),y.shape[2]))
+            bn_tensor_af=mask.new_zeros(*y.shape)
+            start_index=[]
+            ssum=0
+            for i in range(x.shape[0]):
+                start_index.append(ssum)
+                ssum+=index[i]
+            start_index.append(ssum)
+            for i in range(x.shape[0]):
+                bn_tensor_bf[start_index[i]:start_index[i+1]]=y[i,0:index[i]]
+            bn_tensor_bf=self.bn_layer(bn_tensor_bf)
+            for i in range(x.shape[0]):
+                bn_tensor_af[i,0:index[i]]=bn_tensor_bf[start_index[i]:start_index[i+1]]
+            y=bn_tensor_af
+        if self.dropout > 0.001:
+            y = self.dropout_layer(y)
+        if self.relu=='relu':
+            y=torch.nn.functional.relu(y)
+            print('hahah')
+        elif self.relu=='lrelu':
+            y=torch.nn.functional.leaky_relu(y,0.1)
+        return y
+
+#experimental function, untested
+class masked_batchnorm(nn.Module):
+    def __init__(self,feat_dim,epsilon=1e-10):
+        super().__init__()
+        self.alpha=nn.Parameter(torch.ones(feat_dim))
+        self.beta=nn.Parameter(torch.zeros(feat_dim))
+        self.eps=epsilon
+
+    def forward(self,x,mask):
+        '''
+        x: node feat, [batch,node_num,feat_dim]
+        mask: [batch,node_num]
+        '''
+        mask1 = mask.unsqueeze(2)
+        mask_sum = mask.sum()
+        mean = x.sum(dim=(0,1),keepdim=True)/(self.eps+mask_sum)
+        temp = (x - mean)**2
+        temp = temp*mask1
+        var = temp.sum(dim=(0,1),keepdim=True)/(self.eps+mask_sum)
+        rstd = torch.rsqrt(var+self.eps)
+        x=(x-mean)*rstd
+        return ((x*self.alpha) + self.beta)*mask1

models/layers.py

@@ -0,0 +1,280 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['forward_hook', 'Clone', 'Add', 'Cat', 'ReLU', 'GELU', 'Dropout', 'BatchNorm2d', 'Linear', 'MaxPool2d',
+           'AdaptiveAvgPool2d', 'AvgPool2d', 'Conv2d', 'Sequential', 'safe_divide', 'einsum', 'Softmax', 'IndexSelect',
+           'LayerNorm', 'AddEye']
+
+
+def safe_divide(a, b):
+    den = b.clamp(min=1e-9) + b.clamp(max=1e-9)
+    den = den + den.eq(0).type(den.type()) * 1e-9
+    return a / den * b.ne(0).type(b.type())
+
+
+def forward_hook(self, input, output):
+    if type(input[0]) in (list, tuple):
+        self.X = []
+        for i in input[0]:
+            x = i.detach()
+            x.requires_grad = True
+            self.X.append(x)
+    else:
+        self.X = input[0].detach()
+        self.X.requires_grad = True
+
+    self.Y = output
+
+
+def backward_hook(self, grad_input, grad_output):
+    self.grad_input = grad_input
+    self.grad_output = grad_output
+
+
+class RelProp(nn.Module):
+    def __init__(self):
+        super(RelProp, self).__init__()
+        # if not self.training:
+        self.register_forward_hook(forward_hook)
+
+    def gradprop(self, Z, X, S):
+        C = torch.autograd.grad(Z, X, S, retain_graph=True)
+        return C
+
+    def relprop(self, R, alpha):
+        return R
+
+class RelPropSimple(RelProp):
+    def relprop(self, R, alpha):
+        Z = self.forward(self.X)
+        S = safe_divide(R, Z)
+        C = self.gradprop(Z, self.X, S)
+
+        if torch.is_tensor(self.X) == False:
+            outputs = []
+            outputs.append(self.X[0] * C[0])
+            outputs.append(self.X[1] * C[1])
+        else:
+            outputs = self.X * (C[0])
+        return outputs
+
+class AddEye(RelPropSimple):
+    # input of shape B, C, seq_len, seq_len
+    def forward(self, input):
+        return input + torch.eye(input.shape[2]).expand_as(input).to(input.device)
+
+class ReLU(nn.ReLU, RelProp):
+    pass
+
+class GELU(nn.GELU, RelProp):
+    pass
+
+class Softmax(nn.Softmax, RelProp):
+    pass
+
+class LayerNorm(nn.LayerNorm, RelProp):
+    pass
+
+class Dropout(nn.Dropout, RelProp):
+    pass
+
+
+class MaxPool2d(nn.MaxPool2d, RelPropSimple):
+    pass
+
+class LayerNorm(nn.LayerNorm, RelProp):
+    pass
+
+class AdaptiveAvgPool2d(nn.AdaptiveAvgPool2d, RelPropSimple):
+    pass
+
+
+class AvgPool2d(nn.AvgPool2d, RelPropSimple):
+    pass
+
+
+class Add(RelPropSimple):
+    def forward(self, inputs):
+        return torch.add(*inputs)
+
+    def relprop(self, R, alpha):
+        Z = self.forward(self.X)
+        S = safe_divide(R, Z)
+        C = self.gradprop(Z, self.X, S)
+
+        a = self.X[0] * C[0]
+        b = self.X[1] * C[1]
+
+        a_sum = a.sum()
+        b_sum = b.sum()
+
+        a_fact = safe_divide(a_sum.abs(), a_sum.abs() + b_sum.abs()) * R.sum()
+        b_fact = safe_divide(b_sum.abs(), a_sum.abs() + b_sum.abs()) * R.sum()
+
+        a = a * safe_divide(a_fact, a.sum())
+        b = b * safe_divide(b_fact, b.sum())
+
+        outputs = [a, b]
+
+        return outputs
+
+class einsum(RelPropSimple):
+    def __init__(self, equation):
+        super().__init__()
+        self.equation = equation
+    def forward(self, *operands):
+        return torch.einsum(self.equation, *operands)
+
+class IndexSelect(RelProp):
+    def forward(self, inputs, dim, indices):
+        self.__setattr__('dim', dim)
+        self.__setattr__('indices', indices)
+
+        return torch.index_select(inputs, dim, indices)
+
+    def relprop(self, R, alpha):
+        Z = self.forward(self.X, self.dim, self.indices)
+        S = safe_divide(R, Z)
+        C = self.gradprop(Z, self.X, S)
+
+        if torch.is_tensor(self.X) == False:
+            outputs = []
+            outputs.append(self.X[0] * C[0])
+            outputs.append(self.X[1] * C[1])
+        else:
+            outputs = self.X * (C[0])
+        return outputs
+
+
+
+class Clone(RelProp):
+    def forward(self, input, num):
+        self.__setattr__('num', num)
+        outputs = []
+        for _ in range(num):
+            outputs.append(input)
+
+        return outputs
+
+    def relprop(self, R, alpha):
+        Z = []
+        for _ in range(self.num):
+            Z.append(self.X)
+        S = [safe_divide(r, z) for r, z in zip(R, Z)]
+        C = self.gradprop(Z, self.X, S)[0]
+
+        R = self.X * C
+
+        return R
+
+class Cat(RelProp):
+    def forward(self, inputs, dim):
+        self.__setattr__('dim', dim)
+        return torch.cat(inputs, dim)
+
+    def relprop(self, R, alpha):
+        Z = self.forward(self.X, self.dim)
+        S = safe_divide(R, Z)
+        C = self.gradprop(Z, self.X, S)
+
+        outputs = []
+        for x, c in zip(self.X, C):
+            outputs.append(x * c)
+
+        return outputs
+
+
+class Sequential(nn.Sequential):
+    def relprop(self, R, alpha):
+        for m in reversed(self._modules.values()):
+            R = m.relprop(R, alpha)
+        return R
+
+class BatchNorm2d(nn.BatchNorm2d, RelProp):
+    def relprop(self, R, alpha):
+        X = self.X
+        beta = 1 - alpha
+        weight = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3) / (
+            (self.running_var.unsqueeze(0).unsqueeze(2).unsqueeze(3).pow(2) + self.eps).pow(0.5))
+        Z = X * weight + 1e-9
+        S = R / Z
+        Ca = S * weight
+        R = self.X * (Ca)
+        return R
+
+
+class Linear(nn.Linear, RelProp):
+    def relprop(self, R, alpha):
+        beta = alpha - 1
+        pw = torch.clamp(self.weight, min=0)
+        nw = torch.clamp(self.weight, max=0)
+        px = torch.clamp(self.X, min=0)
+        nx = torch.clamp(self.X, max=0)
+
+        def f(w1, w2, x1, x2):
+            Z1 = F.linear(x1, w1)
+            Z2 = F.linear(x2, w2)
+            S1 = safe_divide(R, Z1 + Z2)
+            S2 = safe_divide(R, Z1 + Z2)
+            C1 = x1 * torch.autograd.grad(Z1, x1, S1)[0]
+            C2 = x2 * torch.autograd.grad(Z2, x2, S2)[0]
+
+            return C1 + C2
+
+        activator_relevances = f(pw, nw, px, nx)
+        inhibitor_relevances = f(nw, pw, px, nx)
+
+        R = alpha * activator_relevances - beta * inhibitor_relevances
+
+        return R
+
+
+class Conv2d(nn.Conv2d, RelProp):
+    def gradprop2(self, DY, weight):
+        Z = self.forward(self.X)
+
+        output_padding = self.X.size()[2] - (
+                (Z.size()[2] - 1) * self.stride[0] - 2 * self.padding[0] + self.kernel_size[0])
+
+        return F.conv_transpose2d(DY, weight, stride=self.stride, padding=self.padding, output_padding=output_padding)
+
+    def relprop(self, R, alpha):
+        if self.X.shape[1] == 3:
+            pw = torch.clamp(self.weight, min=0)
+            nw = torch.clamp(self.weight, max=0)
+            X = self.X
+            L = self.X * 0 + \
+                torch.min(torch.min(torch.min(self.X, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3,
+                          keepdim=True)[0]
+            H = self.X * 0 + \
+                torch.max(torch.max(torch.max(self.X, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3,
+                          keepdim=True)[0]
+            Za = torch.conv2d(X, self.weight, bias=None, stride=self.stride, padding=self.padding) - \
+                 torch.conv2d(L, pw, bias=None, stride=self.stride, padding=self.padding) - \
+                 torch.conv2d(H, nw, bias=None, stride=self.stride, padding=self.padding) + 1e-9
+
+            S = R / Za
+            C = X * self.gradprop2(S, self.weight) - L * self.gradprop2(S, pw) - H * self.gradprop2(S, nw)
+            R = C
+        else:
+            beta = alpha - 1
+            pw = torch.clamp(self.weight, min=0)
+            nw = torch.clamp(self.weight, max=0)
+            px = torch.clamp(self.X, min=0)
+            nx = torch.clamp(self.X, max=0)
+
+            def f(w1, w2, x1, x2):
+                Z1 = F.conv2d(x1, w1, bias=None, stride=self.stride, padding=self.padding)
+                Z2 = F.conv2d(x2, w2, bias=None, stride=self.stride, padding=self.padding)
+                S1 = safe_divide(R, Z1)
+                S2 = safe_divide(R, Z2)
+                C1 = x1 * self.gradprop(Z1, x1, S1)[0]
+                C2 = x2 * self.gradprop(Z2, x2, S2)[0]
+                return C1 + C2
+
+            activator_relevances = f(pw, nw, px, nx)
+            inhibitor_relevances = f(nw, pw, px, nx)
+
+            R = alpha * activator_relevances - beta * inhibitor_relevances
+        return R

models/weight_init.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+# -*- coding:UTF-8 -*-
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+
+
+def weight_init(m):
+    '''
+    Usage:
+        model = Model()
+        model.apply(weight_init)
+    '''
+    if isinstance(m, nn.Conv1d):
+        init.normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.Conv2d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.Conv3d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose1d):
+        init.normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose2d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose3d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.BatchNorm1d):
+        init.normal_(m.weight.data, mean=1, std=0.02)
+        init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.BatchNorm2d):
+        init.normal_(m.weight.data, mean=1, std=0.02)
+        init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.BatchNorm3d):
+        init.normal_(m.weight.data, mean=1, std=0.02)
+        init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.Linear):
+        init.xavier_normal_(m.weight.data)
+        init.normal_(m.bias.data)
+    elif isinstance(m, nn.LSTM):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+    elif isinstance(m, nn.LSTMCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+    elif isinstance(m, nn.GRU):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+    elif isinstance(m, nn.GRUCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+
+
+if __name__ == '__main__':
+    pass

option.py

@@ -0,0 +1,41 @@
+###########################################################################
+# Created by: YI ZHENG
+# Email: [email protected]
+# Copyright (c) 2020
+###########################################################################
+
+import os
+import argparse
+import torch
+
+class Options():
+    def __init__(self):
+        parser = argparse.ArgumentParser(description='PyTorch Classification')
+        parser.add_argument('--data_path', type=str, help='path to dataset where images store')
+        parser.add_argument('--train_set', type=str, help='train')
+        parser.add_argument('--val_set', type=str, help='validation')
+        parser.add_argument('--model_path', type=str, help='path to trained model')
+        parser.add_argument('--log_path', type=str, help='path to log files')
+        parser.add_argument('--task_name', type=str, help='task name for naming saved model files and log files')
+        parser.add_argument('--train', action='store_true', default=False, help='train only')
+        parser.add_argument('--test', action='store_true', default=False, help='test only')
+        parser.add_argument('--batch_size', type=int, default=6, help='batch size for origin global image (without downsampling)')
+        parser.add_argument('--log_interval_local', type=int, default=10, help='classification classes')
+        parser.add_argument('--resume', type=str, default="", help='path for model')
+        parser.add_argument('--graphcam', action='store_true', default=False, help='GraphCAM')
+        parser.add_argument('--dataset_metadata_path', type=str, help='Location of the metadata associated with the created dataset: label mapping, splits and so on')
+
+
+        # the parser
+        self.parser = parser
+
+    def parse(self):
+        args = self.parser.parse_args()
+        # default settings for epochs and lr
+
+        args.num_epochs = 120
+        args.lr = 1e-3             
+
+        if args.test:
+            args.num_epochs = 1
+        return args