breast pathology pass!

Author: yifanfeng97

HyperG/conv/__init__.py

@@ -1 +1,3 @@
-from .hyconv import HyConv
+from .hyconv import HyConv
+
+__all__ = ['HyConv']

HyperG/conv/hyconv.py

@@ -49,6 +49,7 @@ def gen_node_ft(self, x: torch.Tensor, H: torch.Tensor):
         return x
 
     def forward(self, x: torch.Tensor, H: torch.Tensor, hyedge_weight=None):
+        assert len(x.shape) == 2, 'the input of HyperConv should be N x C'
         # feature transform
         x = x.matmul(self.theta)
 

HyperG/hyedge/gather_neighbor.py

@@ -76,13 +76,14 @@ def neighbor_distance(x: torch.Tensor, k_nearest, dis_metric=pairwise_euclidean_
     :param k_nearest:
     :return:
     """
-    assert len(x.shape) == 2
+
+    assert len(x.shape) == 2, 'should be a tensor with (N x C) or (B x C x M x N)'
 
     # N x C
     node_num = x.size(0)
     dis_matrix = dis_metric(x)
     _, nn_idx = torch.topk(dis_matrix, k_nearest, dim=1, largest=False)
-    hyedge_idx = torch.arange(node_num).unsqueeze(0).repeat(k_nearest, 1).transpose(1, 0).reshape(-1)
+    hyedge_idx = torch.arange(node_num).to(x.device).unsqueeze(0).repeat(k_nearest, 1).transpose(1, 0).reshape(-1)
     H = torch.stack([nn_idx.reshape(-1), hyedge_idx])
     return H
 

HyperG/models/BaseCNNs.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+import torchvision
+
+
+class ResNetFeature(nn.Module):
+
+    def __init__(self, depth=34, pretrained=True):
+        super().__init__()
+        assert depth in [18, 34, 50, 101, 152]
+
+        if depth == 18:
+            base_model = torchvision.models.resnet18(pretrained=pretrained)
+            self.len_feature = 512
+            self.features = nn.Sequential(*list(base_model.children())[:-2])
+        elif depth == 34:
+            base_model = torchvision.models.resnet34(pretrained=pretrained)
+            self.len_feature = 512
+            self.features = nn.Sequential(*list(base_model.children())[:-2])
+        elif depth == 50:
+            base_model = torchvision.models.resnet50(pretrained=pretrained)
+            self.len_feature = 2048
+            self.features = nn.Sequential(*list(base_model.children())[:-2])
+        elif depth == 101:
+            base_model = torchvision.models.resnet101(pretrained=pretrained)
+            self.len_feature = 2048
+            self.features = nn.Sequential(*list(base_model.children())[:-2])
+        elif depth == 152:
+            base_model = torchvision.models.resnet152(pretrained=pretrained)
+            self.len_feature = 2048
+            self.features = nn.Sequential(*list(base_model.children())[:-2])
+        else:
+            raise NotImplementedError(f'ResNet-{depth} is not implemented!')
+
+    def forward(self, x):
+        x = self.features(x)
+
+        # Attention! No reshape!
+        return x
+
+
+class ResNetClassifier(nn.Module):
+
+    def __init__(self, n_class, len_feature):
+        super().__init__()
+        self.len_feature = len_feature
+        self.classifier = nn.Linear(self.len_feature, n_class)
+
+    def forward(self, x):
+        # -> batch_size x C x N
+        x = x.view(x.size(0), x.size(1), -1)
+
+        # -> batch_size x C
+        x = x.mean(dim=-1)
+
+        x = self.classifier(x)
+        return x

HyperG/models/CNN_HGNN.py

@@ -0,0 +1,45 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from HyperG.conv import HyConv
+from HyperG.hyedge import neighbor_distance
+from HyperG.models import ResNetFeature, ResNetClassifier
+
+
+class ResNet_HGNN(nn.Module):
+    def __init__(self, n_class, depth, k_nearest, hiddens=[512], dropout=0.5, pretrained=True):
+        super().__init__()
+        self.dropout = dropout
+        self.k_nearest = k_nearest
+        self.ft_layers = ResNetFeature(depth=depth, pretrained=pretrained)
+
+        # hypergraph convolution for feature refine
+        self.hyconvs = []
+        dim_in = self.ft_layers.len_feature
+        for h in hiddens:
+            dim_out = h
+            self.hyconvs.append(HyConv(dim_in, dim_out))
+            dim_in = dim_out
+        self.hyconvs = nn.ModuleList(self.hyconvs)
+
+        self.cls_layers = ResNetClassifier(n_class=n_class, len_feature=dim_in)
+
+    def forward(self, x):
+        x = self.ft_layers(x)
+
+        assert x.size(0) == 1, 'when construct hypergraph, only support batch size = 1!'
+        x = x.view(x.size(1), x.size(2) * x.size(3))
+        # -> N x C
+        x = x.permute(1, 0)
+        H = neighbor_distance(x, k_nearest=self.k_nearest)
+        # Hypergraph Convs
+        for hyconv in self.hyconvs:
+            x = hyconv(x, H)
+            x = F.leaky_relu(x, inplace=True)
+            x = F.dropout(x, self.dropout)
+        # N x C -> 1 x C x N
+        x = x.permute(1, 0).unsqueeze(0)
+
+        x = self.cls_layers(x)
+
+        return x

HyperG/models/CNNs.py

@@ -0,0 +1,16 @@
+import torch.nn as nn
+from HyperG.models import ResNetFeature, ResNetClassifier
+
+
+class ResNet(nn.Module):
+    def __init__(self, n_class, depth=34, pretrained=True):
+        super().__init__()
+
+        self.ft_layers = ResNetFeature(depth=depth, pretrained=pretrained)
+        self.cls_layers = ResNetClassifier(n_class=n_class, len_feature=self.ft_layers.len_feature)
+
+    def forward(self, x):
+        x = self.ft_layers(x)
+        x = self.cls_layers(x)
+
+        return x

HyperG/models/HGNN.py

@@ -5,12 +5,12 @@
 
 
 class HGNN(nn.Module):
-    def __init__(self, in_ch, n_class, hidens=[16], dropout=0.5) -> None:
+    def __init__(self, in_ch, n_class, hiddens=[16], dropout=0.5) -> None:
         super().__init__()
         self.dropout = dropout
         _in = in_ch
         self.hyconvs = []
-        for _h in hidens:
+        for _h in hiddens:
             _out = _h
             self.hyconvs.append(HyConv(_in, _out))
             _in = _out
@@ -23,4 +23,4 @@ def forward(self, x, H, hyedge_weight=None):
             x = F.leaky_relu(x, inplace=True)
             x = F.dropout(x, self.dropout)
         x = self.last_hyconv(x, H)
-        return x
+        return F.log_softmax(x, dim=1)

HyperG/models/__init__.py

@@ -1 +1,6 @@
-from .HGNN import HGNN
+from .BaseCNNs import ResNetFeature, ResNetClassifier
+from .CNN_HGNN import ResNet_HGNN
+from .HGNN import HGNN
+from .CNNs import ResNet
+
+__all__ = ['ResNetFeature', 'ResNetClassifier', 'HGNN', 'ResNet_HGNN', 'ResNet']

HyperG/utils/__init__.py

@@ -0,0 +1,11 @@
+import os
+
+
+def check_dir(_dir, make=True):
+    if os.path.exists(_dir):
+        return True
+    else:
+        if make:
+            os.makedirs(_dir)
+            return True
+        return False