dlc2action.model.ms_tcn_modules

   1#
   2# Copyright 2020-present by A. Mathis Group and contributors. All rights reserved.
   3#
   4# This project and all its files are licensed under GNU AGPLv3 or later version. 
   5# A copy is included in dlc2action/LICENSE.AGPL.
   6#
   7import copy
   8
   9import torch
  10import torch.nn as nn
  11import torch.nn.functional as F
  12from dlc2action.model.asformer import AttLayer
  13
  14
  15class Refinement(nn.Module):
  16    """
  17    Refinement module
  18    """
  19
  20    def __init__(
  21        self,
  22        num_layers,
  23        num_f_maps_input,
  24        num_f_maps,
  25        dim,
  26        num_classes,
  27        dropout_rate,
  28        direction,
  29        skip_connections,
  30        attention="none",
  31        block_size=0,
  32    ):
  33        """
  34        Parameters
  35        ----------
  36        num_layers : int
  37            the number of layers
  38        num_f_maps : int
  39            the number of feature maps
  40        dim : int
  41            the number of features in input
  42        num_classes : int
  43            the number of target classes
  44        dropout_rate : float
  45            dropout rate
  46        direction : [None, 'forward', 'backward']
  47            the direction of convolutions; if None, regular convolutions are used
  48        skip_connections : bool
  49            if `True`, skip connections are added
  50        block_size : int, default 0
  51            if not 0, skip connections are added to the prediction generation stage with this interval
  52        """
  53
  54        super(Refinement, self).__init__()
  55        self.block_size = block_size
  56        self.direction = direction
  57        if skip_connections:
  58            self.conv_1x1 = nn.Conv1d(dim + num_f_maps_input, num_f_maps, 1)
  59        else:
  60            self.conv_1x1 = nn.Conv1d(dim, num_f_maps, 1)
  61        self.layers = nn.ModuleList(
  62            [
  63                copy.deepcopy(
  64                    DilatedResidualLayer(
  65                        dilation=2**i,
  66                        in_channels=num_f_maps,
  67                        out_channels=num_f_maps,
  68                        dropout_rate=dropout_rate,
  69                        causal=(direction is not None),
  70                    )
  71                )
  72                for i in range(num_layers)
  73            ]
  74        )
  75        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)
  76        self.attention_layers = nn.ModuleList([])
  77        self.attention = attention
  78        if self.attention == "basic":
  79            self.attention_layers += nn.ModuleList(
  80                [
  81                    nn.Conv1d(num_f_maps, num_f_maps, 3, padding=1),
  82                    nn.ReLU(),
  83                    nn.Conv1d(num_f_maps, num_f_maps, 3, padding=1),
  84                    nn.Sigmoid(),
  85                ]
  86            )
  87
  88    def forward(self, x):
  89        """Forward pass."""
  90        x = self.conv_1x1(x)
  91        out = copy.copy(x)
  92        for l_i, layer in enumerate(self.layers):
  93            out = layer(out, self.direction)
  94            if self.block_size != 0 and (l_i + 1) % self.block_size == 0:
  95                out = out + x
  96                x = copy.copy(out)
  97        if self.attention != "none":
  98            x = copy.copy(out)
  99            for layer in self.attention_layers:
 100                x = layer(x)
 101            out = out * x
 102        out = self.conv_out(out)
 103        return out
 104
 105
 106class Refinement_SE(nn.Module):
 107    """
 108    Refinement module
 109    """
 110
 111    def __init__(
 112        self,
 113        num_layers,
 114        num_f_maps_input,
 115        num_f_maps,
 116        dim,
 117        num_classes,
 118        dropout_rate,
 119        direction,
 120        skip_connections,
 121        len_segment,
 122        block_size=0,
 123    ):
 124        """
 125        Parameters
 126        ----------
 127        num_layers : int
 128            the number of layers
 129        num_f_maps : int
 130            the number of feature maps
 131        dim : int
 132            the number of features in input
 133        num_classes : int
 134            the number of target classes
 135        dropout_rate : float
 136            dropout rate
 137        direction : [None, 'forward', 'backward']
 138            the direction of convolutions; if None, regular convolutions are used
 139        skip_connections : bool
 140            if `True`, skip connections are added
 141        block_size : int, default 0
 142            if not 0, skip connections are added to the prediction generation stage with this interval
 143        """
 144
 145        super().__init__()
 146        self.block_size = block_size
 147        self.direction = direction
 148        if skip_connections:
 149            self.conv_1x1 = nn.Conv1d(dim + num_f_maps_input, num_f_maps, 1)
 150        else:
 151            self.conv_1x1 = nn.Conv1d(dim, num_f_maps, 1)
 152        self.layers = nn.ModuleList(
 153            [
 154                DilatedResidualLayer(
 155                    dilation=2**i,
 156                    in_channels=num_f_maps,
 157                    out_channels=num_f_maps,
 158                    dropout_rate=dropout_rate,
 159                    causal=(direction is not None),
 160                )
 161                for i in range(num_layers)
 162            ]
 163        )
 164        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)
 165        self.fc1_f = nn.ModuleList(
 166            [nn.Linear(num_f_maps, num_f_maps // 2) for _ in range(6)]
 167        )
 168        self.fc2_f = nn.ModuleList(
 169            [nn.Linear(num_f_maps // 2, num_f_maps) for _ in range(6)]
 170        )
 171
 172    def _fc1_f(self, tag):
 173        if tag is None:
 174            for i in range(1, len(self.fc1_f)):
 175                self.fc1_f[i].load_state_dict(self.fc1_f[0].state_dict())
 176            return self.fc1_f[0]
 177        else:
 178            return self.fc1_f[tag]
 179
 180    def _fc2_f(self, tag):
 181        if tag is None:
 182            for i in range(1, len(self.fc2_f)):
 183                self.fc2_f[i].load_state_dict(self.fc2_f[0].state_dict())
 184            return self.fc2_f[0]
 185        else:
 186            return self.fc2_f[tag]
 187
 188    def forward(self, x, tag):
 189        """Forward pass."""
 190        x = self.conv_1x1(x)
 191        scale = torch.mean(x, -1)
 192        scale = self._fc2_f(tag)(F.relu(self._fc1_f(tag)(scale)))
 193        scale = F.sigmoid(scale).unsqueeze(-1)
 194        x = scale * x
 195        out = copy.copy(x)
 196        for l_i, layer in enumerate(self.layers):
 197            out = layer(out, self.direction)
 198            if self.block_size != 0 and (l_i + 1) % self.block_size == 0:
 199                out = out + x
 200                x = copy.copy(out)
 201        out = self.conv_out(out)
 202        return out
 203
 204
 205class RefinementB(Refinement):
 206    """
 207    Bidirectional refinement module
 208    """
 209
 210    def forward(self, x):
 211        """Forward pass."""
 212        x_f = self.conv_1x1(x)
 213        x_b = copy.copy(x_f)
 214        forward = copy.copy(x_f)
 215        backward = copy.copy(x_f)
 216        for i, layer_f in enumerate(self.layers):
 217            forward = layer_f(forward, "forward")
 218            backward = layer_f(backward, "backward")
 219            if self.block_size != 0 and (i + 1) % self.block_size == 0:
 220                forward = forward + x_f
 221                backward = backward + x_b
 222                x_f = copy.copy(forward)
 223                x_b = copy.copy(backward)
 224        out = torch.cat([forward, backward], 1)
 225        out = self.conv_out(out)
 226        return out
 227
 228
 229class SimpleResidualLayer(nn.Module):
 230    """
 231    Basic residual layer
 232    """
 233
 234    def __init__(self, num_f_maps, dropout_rate):
 235        """
 236        Parameters
 237        ----------
 238        in_channels : int
 239            number of input channels
 240        out_channels : int
 241            number of output channels
 242        dropout_rate : float
 243            dropout rate
 244        """
 245
 246        super().__init__()
 247        self.conv_1x1_in = nn.Conv1d(num_f_maps, num_f_maps, 1)
 248        self.conv_1x1 = nn.Conv1d(num_f_maps, num_f_maps, 1)
 249        self.dropout = nn.Dropout(dropout_rate)
 250
 251    def forward(self, x):
 252        """Forward pass."""
 253        out = self.conv_1x1_in(x)
 254        out = F.relu(out)
 255        out = self.conv_1x1(out)
 256        out = self.dropout(out)
 257        return x + out
 258
 259
 260class DilatedResidualLayer(nn.Module):
 261    """
 262    Dilated residual layer
 263    """
 264
 265    def __init__(self, dilation, in_channels, out_channels, dropout_rate, causal):
 266        """
 267        Parameters
 268        ----------
 269        dilation : int
 270            dilation
 271        in_channels : int
 272            number of input channels
 273        out_channels : int
 274            number of output channels
 275        dropout_rate : float
 276            dropout rate
 277        causal : bool
 278            if `True`, causal convolutions are used
 279        """
 280
 281        super(DilatedResidualLayer, self).__init__()
 282        self.padding = dilation * 2
 283        self.causal = causal
 284        if self.causal:
 285            padding = 0
 286        else:
 287            padding = dilation
 288        self.conv_dilated = nn.Conv1d(
 289            in_channels, out_channels, 3, padding=padding, dilation=dilation
 290        )
 291        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
 292        self.dropout = nn.Dropout(dropout_rate)
 293
 294    def forward(self, x, direction=None):
 295        """Forward pass."""
 296        if direction is not None and not self.causal:
 297            raise ValueError("Cannot set direction in a non-causal layer!")
 298        elif direction is None and self.causal:
 299            direction = "forward"
 300        if direction == "forward":
 301            padding = (0, self.padding)
 302        elif direction == "backward":
 303            padding = (self.padding, 0)
 304        elif direction is not None:
 305            raise ValueError(
 306                f"Unrecognized direction: {direction}, please choose from"
 307                f'"backward", "forward" and None'
 308            )
 309        if direction is not None:
 310            out = self.conv_dilated(F.pad(x, padding))
 311        else:
 312            out = self.conv_dilated(x)
 313        out = F.relu(out)
 314        out = self.conv_1x1(out)
 315        out = self.dropout(out)
 316        return x + out
 317
 318
 319class DilatedResidualLayer_SE(nn.Module):
 320    """
 321    Dilated residual layer
 322    """
 323
 324    def __init__(
 325        self, dilation, in_channels, out_channels, dropout_rate, causal, len_segment
 326    ):
 327        """
 328        Parameters
 329        ----------
 330        dilation : int
 331            dilation
 332        in_channels : int
 333            number of input channels
 334        out_channels : int
 335            number of output channels
 336        dropout_rate : float
 337            dropout rate
 338        causal : bool
 339            if `True`, causal convolutions are used
 340        """
 341
 342        super().__init__()
 343        self.padding = dilation * 2
 344        self.causal = causal
 345        if self.causal:
 346            padding = 0
 347        else:
 348            padding = dilation
 349        self.conv_dilated = nn.Conv1d(
 350            in_channels, out_channels, 3, padding=padding, dilation=dilation
 351        )
 352        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
 353        self.dropout = nn.Dropout(dropout_rate)
 354        self.fc1_f = nn.ModuleList(
 355            [nn.Linear(out_channels, out_channels // 2) for _ in range(6)]
 356        )
 357        self.fc2_f = nn.ModuleList(
 358            [nn.Linear(out_channels // 2, out_channels) for _ in range(6)]
 359        )
 360        # self.fc1_t = nn.ModuleList([nn.Linear(len_segment, len_segment // 2) for _ in range(6)])
 361        # self.fc2_t = nn.ModuleList([nn.Linear(len_segment // 2, len_segment) for _ in range(6)])
 362
 363    def _fc1_f(self, tag):
 364        if tag is None:
 365            for i in range(1, len(self.fc1_f)):
 366                self.fc1_f[i].load_state_dict(self.fc1_f[0].state_dict())
 367            return self.fc1_f[0]
 368        else:
 369            return self.fc1_f[tag]
 370
 371    def _fc2_f(self, tag):
 372        if tag is None:
 373            for i in range(1, len(self.fc2_f)):
 374                self.fc2_f[i].load_state_dict(self.fc2_f[0].state_dict())
 375            return self.fc2_f[0]
 376        else:
 377            return self.fc2_f[tag]
 378
 379    def _fc1_t(self, tag):
 380        if tag is None:
 381            for i in range(1, len(self.fc1_t)):
 382                self.fc1_t[i].load_state_dict(self.fc1_t[0].state_dict())
 383            return self.fc1_t[0]
 384        else:
 385            return self.fc1_t[tag]
 386
 387    def _fc2_t(self, tag):
 388        if tag is None:
 389            for i in range(1, len(self.fc2_t)):
 390                self.fc2_t[i].load_state_dict(self.fc2_t[0].state_dict())
 391            return self.fc2_t[0]
 392        else:
 393            return self.fc2_t[tag]
 394
 395    def forward(self, x, direction, tag):
 396        """Forward pass."""
 397        if direction is not None and not self.causal:
 398            raise ValueError("Cannot set direction in a non-causal layer!")
 399        elif direction is None and self.causal:
 400            direction = "forward"
 401        if direction == "forward":
 402            padding = (0, self.padding)
 403        elif direction == "backward":
 404            padding = (self.padding, 0)
 405        elif direction is not None:
 406            raise ValueError(
 407                f"Unrecognized direction: {direction}, please choose from"
 408                f'"backward", "forward" and None'
 409            )
 410        if direction is not None:
 411            out = self.conv_dilated(F.pad(x, padding))
 412        else:
 413            out = self.conv_dilated(x)
 414        out = F.relu(out)
 415        out = self.conv_1x1(out)
 416        out = self.dropout(out)
 417        scale = torch.mean(out, -1)
 418        scale = self._fc2_f(tag)(F.relu(self._fc1_f(tag)(scale)))
 419        scale = F.sigmoid(scale).unsqueeze(-1)
 420        # time_scale = torch.mean(out, 1)
 421        # time_scale = self._fc2_t(tag)(F.relu(self._fc1_t(tag)(time_scale)))
 422        # time_scale = F.sigmoid(time_scale).unsqueeze(1)
 423        out = out * scale  # * time_scale
 424        return x + out
 425
 426
 427class DilatedResidualLayer_SEC(nn.Module):
 428    """
 429    Dilated residual layer
 430    """
 431
 432    def __init__(
 433        self, dilation, in_channels, out_channels, dropout_rate, causal, len_segment
 434    ):
 435        """
 436        Parameters
 437        ----------
 438        dilation : int
 439            dilation
 440        in_channels : int
 441            number of input channels
 442        out_channels : int
 443            number of output channels
 444        dropout_rate : float
 445            dropout rate
 446        causal : bool
 447            if `True`, causal convolutions are used
 448        """
 449
 450        super().__init__()
 451        self.padding = dilation * 2
 452        self.causal = causal
 453        if self.causal:
 454            padding = 0
 455        else:
 456            padding = dilation
 457        self.conv_dilated = nn.Conv1d(
 458            in_channels, out_channels, 3, padding=padding, dilation=dilation
 459        )
 460        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
 461        self.dropout = nn.Dropout(dropout_rate)
 462        self.conv_sq = nn.ModuleList(
 463            [
 464                nn.Conv1d(out_channels, 1, 3, padding=padding, dilation=dilation)
 465                for _ in range(6)
 466            ]
 467        )
 468
 469    def _conv_sq(self, tag):
 470        if tag is None:
 471            for i in range(1, len(self.conv_sq)):
 472                self.conv_sq[i].load_state_dict(self.conv_sq[0].state_dict())
 473            return self.conv_sq[0]
 474        else:
 475            return self.conv_sq[tag]
 476
 477    def forward(self, x, direction, tag):
 478        """Forward pass."""
 479        if direction is not None and not self.causal:
 480            raise ValueError("Cannot set direction in a non-causal layer!")
 481        elif direction is None and self.causal:
 482            direction = "forward"
 483        if direction == "forward":
 484            padding = (0, self.padding)
 485        elif direction == "backward":
 486            padding = (self.padding, 0)
 487        elif direction is not None:
 488            raise ValueError(
 489                f"Unrecognized direction: {direction}, please choose from"
 490                f'"backward", "forward" and None'
 491            )
 492        if direction is not None:
 493            out = self.conv_dilated(F.pad(x, padding))
 494        else:
 495            out = self.conv_dilated(x)
 496        out = F.relu(out)
 497        out = self.conv_1x1(out)
 498        out = self.dropout(out)
 499        scale = torch.sigmoid(self._conv_sq(tag)(out))
 500        out = out * scale
 501        return x + out
 502
 503
 504class DualDilatedResidualLayer(nn.Module):
 505    """
 506    Dual dilated residual layer
 507    """
 508
 509    def __init__(
 510        self,
 511        dilation1,
 512        dilation2,
 513        in_channels,
 514        out_channels,
 515        dropout_rate,
 516        causal,
 517        kernel_size=3,
 518    ):
 519        """
 520        Parameters
 521        ----------
 522        dilation1, dilation2 : int
 523            dilation of one of the blocks
 524        in_channels : int
 525            number of input channels
 526        out_channels : int
 527            number of output channels
 528        dropout_rate : float
 529            dropout rate
 530        causal : bool
 531            if `True`, causal convolutions are used
 532        kernel_size : int, default 3
 533            kernel size
 534        """
 535
 536        super(DualDilatedResidualLayer, self).__init__()
 537        self.causal = causal
 538        self.padding1 = dilation1 * (kernel_size - 1) // 2
 539        self.padding2 = dilation2 * (kernel_size - 1) // 2
 540        if self.causal:
 541            self.padding1 *= 2
 542            self.padding2 *= 2
 543        self.conv_dilated_1 = nn.Conv1d(
 544            in_channels,
 545            out_channels,
 546            kernel_size=kernel_size,
 547            padding=self.padding1,
 548            dilation=dilation1,
 549        )
 550        self.conv_dilated_2 = nn.Conv1d(
 551            in_channels,
 552            out_channels,
 553            kernel_size=kernel_size,
 554            padding=self.padding2,
 555            dilation=dilation2,
 556        )
 557        self.conv_fusion = nn.Conv1d(2 * out_channels, out_channels, 1)
 558        self.dropout = nn.Dropout(dropout_rate)
 559
 560    def forward(self, x, direction=None):
 561        """Forward pass."""
 562        if direction is not None and not self.causal:
 563            raise ValueError("Cannot set direction in a non-causal layer!")
 564        elif direction is None and self.causal:
 565            direction = "forward"
 566        if direction not in ["backward", "forward", None]:
 567            raise ValueError(
 568                f"Unrecognized direction: {direction}, please choose from"
 569                f'"backward", "forward" and None'
 570            )
 571        out_1 = self.conv_dilated_1(x)
 572        out_2 = self.conv_dilated_2(x)
 573        if direction == "forward":
 574            out_1 = out_1[:, :, : -self.padding1]
 575            out_2 = out_2[:, :, : -self.padding2]
 576        elif direction == "backward":
 577            out_1 = out_1[:, :, self.padding1 :]
 578            out_2 = out_2[:, :, self.padding2 :]
 579
 580        out = self.conv_fusion(
 581            torch.cat(
 582                [
 583                    out_1,
 584                    out_2,
 585                ],
 586                1,
 587            )
 588        )
 589        out = F.relu(out)
 590        out = self.dropout(out)
 591        out = out + x
 592        return out
 593
 594
 595class MultiDilatedTCN(nn.Module):
 596    """
 597    Multiple prediction generation stages in parallel
 598    """
 599
 600    def __init__(
 601        self,
 602        num_layers,
 603        num_f_maps,
 604        dims,
 605        direction,
 606        block_size=5,
 607        kernel_size=3,
 608        rare_dilations=False,
 609    ):
 610        super(MultiDilatedTCN, self).__init__()
 611        self.PGs = nn.ModuleList(
 612            [
 613                DilatedTCN(
 614                    num_layers,
 615                    num_f_maps,
 616                    dim,
 617                    direction,
 618                    block_size,
 619                    kernel_size,
 620                    rare_dilations,
 621                )
 622                for dim in dims
 623            ]
 624        )
 625        self.conv_out = nn.Conv1d(num_f_maps * len(dims), num_f_maps, 1)
 626
 627    def forward(self, x, tag=None):
 628        """Forward pass."""
 629        out = []
 630        for arr, PG in zip(x, self.PGs):
 631            out.append(PG(arr))
 632        out = torch.cat(out, 1)
 633        out = self.conv_out(out)
 634        return out
 635
 636
 637class DilatedTCN(nn.Module):
 638    """
 639    Prediction generation stage
 640    """
 641
 642    def __init__(
 643        self,
 644        num_layers,
 645        num_f_maps,
 646        dim,
 647        direction,
 648        num_bp=None,
 649        block_size=5,
 650        kernel_size=3,
 651        rare_dilations=False,
 652        attention="none",
 653        multihead=False,
 654    ):
 655        """
 656        Parameters
 657        ----------
 658        num_layers : int
 659            number of layers
 660        num_f_maps : int
 661            number of feature maps
 662        dim : int
 663            number of features in input
 664        direction : [None, 'forward', 'backward']
 665            the direction of convolutions; if None, regular convolutions are used
 666        block_size : int, default 0
 667            if not 0, skip connections are added to the prediction generation stage with this interval
 668        kernel_size : int, default 3
 669            kernel size
 670        rare_dilations : bool, default False
 671            if `False`, dilation increases every layer, otherwise every second layer
 672        """
 673
 674        super().__init__()
 675        self.num_layers = num_layers
 676        self.block_size = block_size
 677        self.direction = direction
 678        self.conv_1x1_in = nn.Conv1d(dim, num_f_maps, 1)
 679        pars = {
 680            "in_channels": num_f_maps,
 681            "out_channels": num_f_maps,
 682            "dropout_rate": 0.5,
 683            "causal": (direction is not None),
 684            "kernel_size": kernel_size,
 685        }
 686        module = DualDilatedResidualLayer
 687        if not rare_dilations:
 688            self.layers = nn.ModuleList(
 689                [
 690                    module(dilation1=2 ** (num_layers - 1 - i), dilation2=2**i, **pars)
 691                    for i in range(num_layers)
 692                ]
 693            )
 694        else:
 695            self.layers = nn.ModuleList(
 696                [
 697                    module(
 698                        dilation1=2 ** (num_layers // 2 - 1 - i // 2),
 699                        dilation2=2 ** (i // 2),
 700                        **pars,
 701                    )
 702                    for i in range(num_layers)
 703                ]
 704            )
 705        self.attention_layers = nn.ModuleList([])
 706        self.attention = attention
 707        if self.attention in ["basic"]:
 708            self.attention_layers += nn.ModuleList(
 709                [
 710                    nn.Conv1d(num_f_maps, num_f_maps, 3, padding=1),
 711                    nn.ReLU(),
 712                    nn.Conv1d(num_f_maps, num_f_maps, 3, padding=1),
 713                    nn.Sigmoid(),
 714                ]
 715            )
 716        elif isinstance(self.attention, str) and self.attention != "none":
 717            self.attention_layers = AttLayer(
 718                num_f_maps,
 719                num_f_maps,
 720                num_f_maps,
 721                4,
 722                4,
 723                2,
 724                64,
 725                att_type=self.attention,
 726                stage="encoder",
 727            )
 728        self.multihead = multihead
 729
 730    def forward(self, x, tag=None):
 731        """Forward pass."""
 732        x = self.conv_1x1_in(x)
 733        f = copy.copy(x)
 734        for i, layer in enumerate(self.layers):
 735            f = layer(f, self.direction)
 736            if self.block_size != 0 and (i + 1) % self.block_size == 0:
 737                f = f + x
 738                x = copy.copy(f)
 739        if isinstance(self.attention, str) and self.attention != "none":
 740            x = copy.copy(f)
 741            if self.attention != "basic":
 742                f = self.attention_layers(x)
 743            elif not self.multihead:
 744                for layer in self.attention_layers:
 745                    x = layer(x)
 746                f = f * x
 747            else:
 748                outputs = []
 749                for layers in self.attention_layers:
 750                    y = copy.copy(x)
 751                    for layer in layers:
 752                        y = layer(y)
 753                    outputs.append(copy.copy(f * y))
 754                outputs = torch.cat(outputs, dim=1)
 755                f = self.conv_att(self.dropout(outputs))
 756        return f
 757
 758
 759class DilatedTCNB(nn.Module):
 760    """
 761    Bidirectional prediction generation stage
 762    """
 763
 764    def __init__(
 765        self,
 766        num_layers,
 767        num_f_maps,
 768        dim,
 769        block_size=5,
 770        kernel_size=3,
 771        rare_dilations=False,
 772    ):
 773        """
 774        Parameters
 775        ----------
 776        num_layers : int
 777            number of layers
 778        num_f_maps : int
 779            number of feature maps
 780        dim : int
 781            number of features in input
 782        block_size : int, default 0
 783            if not 0, skip connections are added to the prediction generation stage with this interval
 784        kernel_size : int, default 3
 785            kernel size
 786        rare_dilations : bool, default False
 787            if `False`, dilation increases every layer, otherwise every second layer
 788        """
 789
 790        super().__init__()
 791        self.num_layers = num_layers
 792        self.block_size = block_size
 793        self.conv_1x1_in = nn.Conv1d(dim, num_f_maps, 1)
 794        self.conv_1x1_out = nn.Conv1d(num_f_maps * 2, num_f_maps, 1)
 795        if not rare_dilations:
 796            self.layers = nn.ModuleList(
 797                [
 798                    DualDilatedResidualLayer(
 799                        dilation1=2 ** (num_layers - 1 - i),
 800                        dilation2=2**i,
 801                        in_channels=num_f_maps,
 802                        out_channels=num_f_maps,
 803                        dropout_rate=0.5,
 804                        causal=True,
 805                        kernel_size=kernel_size,
 806                    )
 807                    for i in range(num_layers)
 808                ]
 809            )
 810        else:
 811            self.layers = nn.ModuleList(
 812                [
 813                    DualDilatedResidualLayer(
 814                        dilation1=2 ** (num_layers // 2 - 1 - i // 2),
 815                        dilation2=2 ** (i // 2),
 816                        in_channels=num_f_maps,
 817                        out_channels=num_f_maps,
 818                        dropout_rate=0.5,
 819                        causal=True,
 820                        kernel_size=kernel_size,
 821                    )
 822                    for i in range(num_layers)
 823                ]
 824            )
 825
 826    def forward(self, x, tag=None):
 827        """Forward pass."""
 828        x_f = self.conv_1x1_in(x)
 829        x_b = copy.copy(x_f)
 830        forward = copy.copy(x_f)
 831        backward = copy.copy(x_f)
 832        for i, layer_f in enumerate(self.layers):
 833            forward = layer_f(forward, "forward")
 834            backward = layer_f(backward, "backward")
 835            if self.block_size != 0 and (i + 1) % self.block_size == 0:
 836                forward = forward + x_f
 837                backward = backward + x_b
 838                x_f = copy.copy(forward)
 839                x_b = copy.copy(backward)
 840        out = torch.cat([forward, backward], 1)
 841        out = self.conv_1x1_out(out)
 842        return out
 843
 844
 845class SpatialFeatures(nn.Module):
 846    """
 847    Spatial features extraction stage
 848    """
 849
 850    def __init__(
 851        self,
 852        num_layers,
 853        num_f_maps,
 854        dim,
 855        block_size=5,
 856        graph_edges=None,
 857        num_nodes=None,
 858        denom: int = 8,
 859    ):
 860        """
 861        Parameters
 862        ----------
 863        num_layers : int
 864            number of layers
 865        num_f_maps : int
 866            number of feature maps
 867        dim : int
 868            number of features in input
 869        block_size : int, default 5
 870            if not 0, skip connections are added to the prediction generation stage with this interval
 871        """
 872
 873        super().__init__()
 874        self.num_nodes = num_nodes
 875        if graph_edges is None:
 876            module = SimpleResidualLayer
 877            self.graph = False
 878            pars = {"num_f_maps": num_f_maps, "dropout_rate": 0.5}
 879            self.conv_1x1_in = nn.Conv1d(dim, num_f_maps, 1)
 880        else:
 881            raise NotImplementedError("Graph not implemented")
 882        
 883        self.num_layers = num_layers
 884        self.block_size = block_size
 885        self.layers = nn.ModuleList([module(**pars) for _ in range(num_layers)])
 886
 887    def forward(self, x):
 888        """Forward pass."""
 889        if self.graph:
 890            B, _, L = x.shape
 891            x = x.transpose(-1, -2)
 892            x = x.reshape((-1, x.shape[-1]))
 893            x = x.reshape((x.shape[0], self.num_nodes, -1))
 894            x = x.transpose(-1, -2)
 895        x = self.conv_1x1_in(x)
 896        f = copy.copy(x)
 897        for i, layer in enumerate(self.layers):
 898            f = layer(f)
 899            if self.block_size != 0 and (i + 1) % self.block_size == 0:
 900                f = f + x
 901                x = copy.copy(f)
 902        if self.graph:
 903            f = f.reshape((B, L, -1))
 904            f = f.transpose(-1, -2)
 905            f = self.conv_1x1_out(f)
 906        return f
 907
 908
 909class MSRefinement(nn.Module):
 910    """
 911    Refinement stage
 912    """
 913
 914    def __init__(
 915        self,
 916        num_layers_R,
 917        num_R,
 918        num_f_maps_input,
 919        num_f_maps,
 920        num_classes,
 921        dropout_rate,
 922        exclusive,
 923        skip_connections,
 924        direction,
 925        block_size=0,
 926        num_heads=1,
 927        attention="none",
 928    ):
 929        """
 930        Parameters
 931        ----------
 932        num_layers_R : int
 933            number of layers in refinement modules
 934        num_R : int
 935            number of refinement modules
 936        num_f_maps : int
 937            number of feature maps
 938        num_classes : int
 939            number of target classes
 940        dropout_rate : float
 941            dropout rate
 942        exclusive : bool
 943            set `False` for multi-label classification
 944        skip_connections : bool
 945            if `True`, skip connections are added
 946        direction : [None, 'bidirectional', 'forward', 'backward']
 947            the direction of convolutions; if None, regular convolutions are used
 948        block_size : int, default 0
 949            if not 0, skip connections are added to the prediction generation stage with this interval
 950        num_heads : int, default 1
 951            number of parallel refinement stages
 952        """
 953
 954        super().__init__()
 955        self.skip_connections = skip_connections
 956        self.num_heads = num_heads
 957        if exclusive:
 958            self.nl = lambda x: F.softmax(x, dim=1)
 959        else:
 960            self.nl = lambda x: torch.sigmoid(x)
 961        if direction == "bidirectional":
 962            refinement_module = RefinementB
 963        else:
 964            refinement_module = Refinement
 965        self.Rs = nn.ModuleList(
 966            [
 967                nn.ModuleList(
 968                    [
 969                        refinement_module(
 970                            num_layers=num_layers_R,
 971                            num_f_maps=num_f_maps,
 972                            num_f_maps_input=num_f_maps_input,
 973                            dim=num_classes,
 974                            num_classes=num_classes,
 975                            dropout_rate=dropout_rate,
 976                            direction=direction,
 977                            skip_connections=skip_connections,
 978                            block_size=block_size,
 979                            attention=attention,
 980                        )
 981                        for s in range(num_R)
 982                    ]
 983                )
 984                for _ in range(self.num_heads)
 985            ]
 986        )
 987        self.conv_out = nn.ModuleList(
 988            [nn.Conv1d(num_f_maps_input, num_classes, 1) for _ in range(self.num_heads)]
 989        )
 990        if self.num_heads == 1:
 991            self.Rs = self.Rs[0]
 992            self.conv_out = self.conv_out[0]
 993
 994    def _Rs(self, tag):
 995        if self.num_heads == 1:
 996            return self.Rs
 997        if tag is None:
 998            tag = 0
 999            for i in range(1, self.num_heads):
1000                self.Rs[i].load_state_dict(self.Rs[0].state_dict())
1001        return self.Rs[tag]
1002
1003    def _conv_out(self, tag):
1004        if self.num_heads == 1:
1005            return self.conv_out
1006        if tag is None:
1007            tag = 0
1008            for i in range(1, self.num_heads):
1009                self.conv_out[i].load_state_dict(self.conv_out[0].state_dict())
1010        return self.conv_out[tag]
1011
1012    def forward(self, x, tag=None):
1013        """Forward pass."""
1014        if tag is not None:
1015            tag = tag[0]
1016        out = self._conv_out(tag)(x)
1017        outputs = out.unsqueeze(0)
1018        for R in self._Rs(tag):
1019            if self.skip_connections:
1020                out = R(torch.cat([self.nl(out), x], axis=1))
1021                # out = R(torch.cat([out, x], axis=1))
1022            else:
1023                out = R(self.nl(out))
1024                # out = R(out)
1025            outputs = torch.cat((outputs, out.unsqueeze(0)), dim=0)
1026
1027        return outputs
1028
1029
1030class MSRefinementShared(nn.Module):
1031    """
1032    Refinement stage with shared weights across modules
1033    """
1034
1035    def __init__(
1036        self,
1037        num_layers_R,
1038        num_R,
1039        num_f_maps_input,
1040        num_f_maps,
1041        num_classes,
1042        dropout_rate,
1043        exclusive,
1044        skip_connections,
1045        direction,
1046        block_size=0,
1047        num_heads=1,
1048        attention="none",
1049    ):
1050        """
1051        Parameters
1052        ----------
1053        num_layers_R : int
1054            number of layers in refinement modules
1055        num_R : int
1056            number of refinement modules
1057        num_f_maps : int
1058            number of feature maps
1059        num_classes : int
1060            number of target classes
1061        dropout_rate : float
1062            dropout rate
1063        exclusive : bool
1064            set `False` for multi-label classification
1065        skip_connections : bool
1066            if `True`, skip connections are added
1067        direction : [None, 'bidirectional', 'forward', 'backward']
1068            the direction of convolutions; if None, regular convolutions are used
1069        block_size : int, default 0
1070            if not 0, skip connections are added to the prediction generation stage with this interval
1071        num_heads : int, default 1
1072            number of parallel refinement stages
1073        """
1074
1075        super().__init__()
1076        if exclusive:
1077            self.nl = lambda x: F.softmax(x, dim=1)
1078        else:
1079            self.nl = lambda x: torch.sigmoid(x)
1080        if direction == "bidirectional":
1081            refinement_module = RefinementB
1082        else:
1083            refinement_module = Refinement
1084        self.num_heads = num_heads
1085        self.R = nn.ModuleList(
1086            [
1087                refinement_module(
1088                    num_layers=num_layers_R,
1089                    num_f_maps_input=num_f_maps_input,
1090                    num_f_maps=num_f_maps,
1091                    dim=num_classes,
1092                    num_classes=num_classes,
1093                    dropout_rate=dropout_rate,
1094                    direction=direction,
1095                    skip_connections=skip_connections,
1096                    block_size=block_size,
1097                    attention=attention,
1098                )
1099                for _ in range(self.num_heads)
1100            ]
1101        )
1102        self.num_R = num_R
1103        self.conv_out = nn.ModuleList(
1104            [nn.Conv1d(num_f_maps_input, num_classes, 1) for _ in range(self.num_heads)]
1105        )
1106        self.skip_connections = skip_connections
1107        if self.num_heads == 1:
1108            self.R = self.R[0]
1109            self.conv_out = self.conv_out[0]
1110
1111    def _R(self, tag):
1112        if self.num_heads == 1:
1113            return self.R
1114        if tag is None:
1115            tag = 0
1116            for i in range(1, self.num_heads):
1117                self.R[i].load_state_dict(self.R[0].state_dict())
1118        return self.R[tag]
1119
1120    def _conv_out(self, tag):
1121        if self.num_heads == 1:
1122            return self.conv_out
1123        if tag is None:
1124            tag = 0
1125            for i in range(1, self.num_heads):
1126                self.conv_out[i].load_state_dict(self.conv_out[0].state_dict())
1127        return self.conv_out[tag]
1128
1129    def forward(self, x, tag=None):
1130        """Forward pass."""
1131        if tag is not None:
1132            tag = tag[0]
1133        out = self._conv_out(tag)(x)
1134        outputs = out.unsqueeze(0)
1135        for _ in range(self.num_R):
1136            if self.skip_connections:
1137                # out = self._R(tag)(torch.cat([self.nl(out), x], axis=1))
1138                out = self._R(tag)(torch.cat([out, x], axis=1))
1139            else:
1140                # out = self._R(tag)(self.nl(out))
1141                out = self._R(tag)(out)
1142            outputs = torch.cat((outputs, out.unsqueeze(0)), dim=0)
1143
1144        return outputs
1145
1146
1147class MSRefinementAttention(nn.Module):
1148    """
1149    Refinement stage
1150    """
1151
1152    def __init__(
1153        self,
1154        num_layers_R,
1155        num_R,
1156        num_f_maps_input,
1157        num_f_maps,
1158        num_classes,
1159        dropout_rate,
1160        exclusive,
1161        skip_connections,
1162        len_segment,
1163        block_size=0,
1164    ):
1165        """
1166        Parameters
1167        ----------
1168        num_layers_R : int
1169            number of layers in refinement modules
1170        num_R : int
1171            number of refinement modules
1172        num_f_maps : int
1173            number of feature maps
1174        num_classes : int
1175            number of target classes
1176        dropout_rate : float
1177            dropout rate
1178        exclusive : bool
1179            set `False` for multi-label classification
1180        skip_connections : bool
1181            if `True`, skip connections are added
1182        direction : [None, 'bidirectional', 'forward', 'backward']
1183            the direction of convolutions; if None, regular convolutions are used
1184        block_size : int, default 0
1185            if not 0, skip connections are added to the prediction generation stage with this interval
1186        num_heads : int, default 1
1187            number of parallel refinement stages
1188        """
1189
1190        super().__init__()
1191        self.skip_connections = skip_connections
1192        if exclusive:
1193            self.nl = lambda x: F.softmax(x, dim=1)
1194        else:
1195            self.nl = lambda x: torch.sigmoid(x)
1196        refinement_module = Refinement_SE
1197        self.Rs = nn.ModuleList(
1198            [
1199                refinement_module(
1200                    num_layers=num_layers_R,
1201                    num_f_maps=num_f_maps,
1202                    num_f_maps_input=num_f_maps_input,
1203                    dim=num_classes,
1204                    num_classes=num_classes,
1205                    dropout_rate=dropout_rate,
1206                    direction=None,
1207                    skip_connections=skip_connections,
1208                    block_size=block_size,
1209                    len_segment=len_segment,
1210                )
1211                for s in range(num_R)
1212            ]
1213        )
1214        self.conv_out = nn.Conv1d(num_f_maps_input, num_classes, 1)
1215
1216    def forward(self, x, tag=None):
1217        """Forward pass."""
1218        if tag is not None:
1219            tag = tag[0]
1220        out = self.conv_out(x)
1221        outputs = out.unsqueeze(0)
1222        for R in self.Rs:
1223            if self.skip_connections:
1224                out = R(torch.cat([self.nl(out), x], axis=1), tag)
1225                # out = R(torch.cat([out, x], axis=1), tag)
1226            else:
1227                out = R(self.nl(out), tag)
1228                # out = R(out, tag)
1229            outputs = torch.cat((outputs, out.unsqueeze(0)), dim=0)
1230
1231        return outputs
1232
1233
1234class DilatedTCNC(nn.Module):
1235    def __init__(
1236        self,
1237        num_f_maps,
1238        num_layers_PG,
1239        len_segment,
1240        block_size_prediction=5,
1241        kernel_size_prediction=3,
1242        direction_PG=None,
1243    ):
1244        super(DilatedTCNC, self).__init__()
1245        if direction_PG == "bidirectional":
1246            self.PG_S = DilatedTCNB(
1247                num_layers=num_layers_PG,
1248                num_f_maps=num_f_maps,
1249                dim=num_f_maps,
1250                block_size=block_size_prediction,
1251            )
1252            self.PG_T = DilatedTCNB(
1253                num_layers=num_layers_PG,
1254                num_f_maps=len_segment,
1255                dim=len_segment,
1256                block_size=block_size_prediction,
1257            )
1258        else:
1259            self.PG_S = DilatedTCN(
1260                num_layers=num_layers_PG,
1261                num_f_maps=num_f_maps,
1262                dim=num_f_maps,
1263                direction=direction_PG,
1264                block_size=block_size_prediction,
1265                kernel_size=kernel_size_prediction,
1266            )
1267            self.PG_T = DilatedTCN(
1268                num_layers=num_layers_PG,
1269                num_f_maps=len_segment,
1270                dim=len_segment,
1271                direction=direction_PG,
1272                block_size=block_size_prediction,
1273                kernel_size=kernel_size_prediction,
1274            )
1275
1276    def forward(self, x):
1277        """Forward pass."""
1278        x = self.PG_S(x)
1279        x = torch.transpose(x, 1, 2)
1280        x = self.PG_T(x)
1281        x = torch.transpose(x, 1, 2)
1282        return x