catalyst_example/meshnetme.py at main · neuroneural/catalyst_example · 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
142
143
import torch
import torch.nn as nn
import torch.nn.functional as F
from meshnet_gn import set_channel_num, construct_layer
from torch.utils.checkpoint import checkpoint_sequential
import json


class MeshNetME(nn.Module):
    def __init__(self, in_channels, n_classes, channels, config_file):
        super().__init__()
        self.num_experts = len(config_file)

        # Initialize the towers
        towers = {}
        for i, config in enumerate(config_file):
            with open(config, "r") as f:
                config = set_channel_num(
                    json.load(f), in_channels, n_classes, channels
                )
            layers = [
                construct_layer(
                    dropout_p=config["dropout_p"],
                    bnorm=config["bnorm"],
                    gelu=config["gelu"],
                    **block_kwargs,
                )
                for block_kwargs in config["layers"]
            ]
            layers[-1] = layers[-1][0]
            towers[f"tower_{i}"] = nn.Sequential(*layers)
        self.towers = nn.ModuleDict(towers)

    def forward(self, x):
        # Compute the outputs of each expert
        # option 1
        # tower_sum = sum([tower(x) for tower in self.towers.values()])
        # memory efficient option 2
        # tower_sum = next(iter(self.towers.values()))(x)
        # # Add the outputs of the remaining towers
        # for tower in list(self.towers.values())[1:]:
        #     tower_sum += tower(x)

        tower_outputs = [tower(x) for tower in self.towers.values()]
        tower_sum = torch.sum(torch.stack(tower_outputs, dim=1), dim=1)

        return tower_sum


class MeshNetME_(nn.Module):
    def __init__(self, in_channels, n_classes, channels, config_files):
        super().__init__()
        self.num_experts = len(config_files)

        # Initialize the towers
        towers = {}
        for i, config_file in enumerate(config_files):
            with open(config_file, "r") as f:
                config = set_channel_num(
                    json.load(f), in_channels, n_classes, channels
                )
            layers = [
                construct_layer(
                    dropout_p=config["dropout_p"],
                    bnorm=config["bnorm"],
                    gelu=config["gelu"],
                    **block_kwargs,
                )
                for block_kwargs in config["layers"]
            ]
            del layers[-1]
            towers[f"tower_{i}"] = nn.Sequential(*layers)
        self.towers = nn.ModuleDict(towers)
        layers = [
            nn.GroupNorm(
                num_groups=channels,
                num_channels=channels,
            ),
            nn.Conv3d(channels, n_classes, kernel_size=1),
        ]
        self._last_layer = nn.Sequential(*layers)

    def forward(self, x):
        # Compute the outputs of each expert
        # option 1
        # tower_sum = sum([tower(x) for tower in self.towers.values()])
        # memory efficient option 2
        # tower_sum = next(iter(self.towers.values()))(x)
        # # Add the outputs of the remaining towers
        # for tower in list(self.towers.values())[1:]:
        #     tower_sum += tower(x)

        tower_outputs = [tower(x) for tower in self.towers.values()]
        tower_sum = torch.sum(torch.stack(tower_outputs, dim=1), dim=1)
        # Apply the final convolution
        output = self._last_layer(tower_sum)
        return output


class CheckpointMixin:
    def train_forward(self, x, model):
        y = x
        y.requires_grad_()
        y = checkpoint_sequential(
            model, len(model), y, preserve_rng_state=False
        )
        return y

    def eval_forward(self, x, model):
        """Forward pass"""
        model.eval()
        with torch.inference_mode():
            x = model(x)
        return x

    def forward(self, x):
        if self.training:
            # tower_outputs = [
            #     self.train_forward(x, tower) for tower in self.towers.values()
            # ]
            tower_sum = self.train_forward(x, next(iter(self.towers.values())))
            # Add the outputs of the remaining towers
            for tower in list(self.towers.values())[1:]:
                tower_sum += self.train_forward(x, tower)
        else:
            # tower_outputs = [
            #     self.eval_forward(x, tower) for tower in self.towers.values()
            # ]
            with torch.inference_mode():
                tower_sum = self.eval_forward(
                    x, next(iter(self.towers.values()))
                )
                # Add the outputs of the remaining towers
                for tower in list(self.towers.values())[1:]:
                    tower_sum += self.eval_forward(x, tower)

        # tower_sum = torch.sum(torch.stack(tower_outputs, dim=1), dim=1)

        return tower_sum


class MeshNetME_checkpoint(CheckpointMixin, MeshNetME):
    pass