HyperCore/hypercore/models/lorentz_resnet.py at main · Graph-and-Geometric-Learning/HyperCore · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import torch.nn as nn
import torch
from ..nn.conv import *
from ..manifolds import Lorentz

class Lorentz_ResNet(nn.Module):
    """ Implementation of ResNet models on manifolds. """
    def __init__(
        self,
        num_blocks,
        manifold_in:Lorentz,
        manifold_hidden:Lorentz,
        manifold_out:Lorentz,
        img_dim=[3,32,32],
        embed_dim=512,
        num_classes=100,
        bias=True,
        remove_linear=False,
    ):
        super(Lorentz_ResNet, self).__init__()

        self.img_dim = img_dim[0]
        self.in_channels = 64
        self.conv3_dim = 128
        self.conv4_dim = 256
        self.embed_dim = embed_dim
        self.bias = bias

        self.c = torch.Tensor([0.1]).cpu()
        self.manifold_in = manifold_in
        self.manifold_hidden = manifold_hidden
        self.manifold_out = manifold_out

        self.conv1 = self._get_inConv()
        self.conv2_x = self._make_layer(out_channels=self.in_channels, num_blocks=num_blocks[0], stride=1)
        self.conv3_x = self._make_layer(out_channels=self.conv3_dim, num_blocks=num_blocks[1], stride=2)
        self.conv4_x = self._make_layer(out_channels=self.conv4_dim, num_blocks=num_blocks[2], stride=2)
        self.conv5_x = self._make_layer(out_channels=self.embed_dim, num_blocks=num_blocks[3], stride=2)
        self.avg_pool = self._get_GlobalAveragePooling()

        if remove_linear:
            self.predictor = None
        else:
            self.predictor = self._get_predictor(self.embed_dim*LorentzResidualBlock.expansion, num_classes)

    def forward(self, x):
        out = self.conv1(x)

        out_1 = self.conv2_x(out)
        out_2 = self.conv3_x(out_1)
        out_3 = self.conv4_x(out_2)
        out_4 = self.conv5_x(out_3)
        out = self.avg_pool(out_4)
        out = out.view(out.size(0), -1)

        if self.predictor is not None:
            out = self.predictor(out)

        return out

    def _make_layer(self, out_channels, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []

        for stride in strides:
            layers.append(
                LorentzResidualBlock(
                    self.manifold_hidden,
                    in_channels=self.in_channels,
                    out_channels=out_channels,
                    stride=stride,
                )
            )
            self.in_channels = out_channels * LorentzResidualBlock.expansion + 1

        return nn.Sequential(*layers)

    def _get_inConv(self):
        return LorentzInputBlock(
            self.manifold_in,
            self.img_dim,
            self.in_channels,
            self.bias,
            self.manifold_hidden
        )

    def _get_predictor(self, in_features, num_classes):
        return LorentzMLR(self.manifold_out, in_features+1, num_classes)

    def _get_GlobalAveragePooling(self):
        return LorentzGlobalAvgPool2d(manifold_in=self.manifold_hidden, keep_dim=True, manifold_out=self.manifold_out)

def Lorentz_resnet18(manifold_in:Lorentz, manifold_hidden=None, manifold_out=None, **kwargs):
    if manifold_out is None:
        if manifold_hidden:
            manifold_out = manifold_hidden
        else:
            manifold_hidden = manifold_in
            manifold_out = manifold_in
    model = Lorentz_ResNet([2, 2, 2, 2], manifold_in, manifold_hidden, manifold_out, **kwargs)
    return model

def Lorentz_resnet34(manifold_in:Lorentz, manifold_hidden=None, manifold_out=None, **kwargs):
    if manifold_out is None:
        if manifold_hidden:
            manifold_out = manifold_hidden
        else:
            manifold_hidden = manifold_in
            manifold_out = manifold_in
    model = Lorentz_ResNet([3, 4, 6, 3], manifold_in, manifold_hidden, manifold_out, **kwargs)
    return model

def Lorentz_resnet50(manifold_in:Lorentz, manifold_hidden=None, manifold_out=None, **kwargs):
    if manifold_out is None:
        if manifold_hidden:
            manifold_out = manifold_hidden
        else:
            manifold_hidden = manifold_in
            manifold_out = manifold_in
    model = Lorentz_ResNet([3, 4, 6, 3], manifold_in, manifold_hidden, manifold_out, **kwargs)
    return model

def Lorentz_resnet101(manifold_in:Lorentz, manifold_hidden=None, manifold_out=None, **kwargs):
    if manifold_out is None:
        if manifold_hidden:
            manifold_out = manifold_hidden
        else:
            manifold_hidden = manifold_in
            manifold_out = manifold_in
    model = Lorentz_ResNet([3, 4, 23, 3], manifold_in, manifold_hidden, manifold_out, **kwargs)
    return model

def Lorentz_resnet152(manifold_in:Lorentz, manifold_hidden=None, manifold_out=None, **kwargs):
    if manifold_out is None:
        if manifold_hidden:
            manifold_out = manifold_hidden
        else:
            manifold_hidden = manifold_in
            manifold_out = manifold_in
    model = Lorentz_ResNet([3, 8, 36, 3], manifold_in, manifold_hidden, manifold_out, **kwargs)
    return model