Class LobTransformer

models/LobTransformer/lobtransformer.py:6–124

class LobTransformer(pl.LightningModule):
    def __init__(self, lighten):
        super().__init__()
        self.name = "lobtransformer"
        if lighten:
            self.name += "-lighten"

        hidden = 32 if not lighten else 16
        d_model = hidden * 2 * 3
        nhead = 8 if not lighten else 4
        num_layers = 2 if not lighten else 1

        # Convolution blocks.
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1, out_channels=hidden, kernel_size=(1, 2), stride=(1, 2)
            ),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
            nn.Conv2d(in_channels=hidden, out_channels=hidden, kernel_size=(4, 1)),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
            nn.Conv2d(in_channels=hidden, out_channels=hidden, kernel_size=(4, 1)),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                in_channels=hidden, out_channels=hidden, kernel_size=(1, 2), stride=(1, 2)
            ),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
            nn.Conv2d(in_channels=hidden, out_channels=hidden, kernel_size=(4, 1)),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
            nn.Conv2d(in_channels=hidden, out_channels=hidden, kernel_size=(4, 1)),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
        )

        if lighten:
            conv3_kernel_size = 5
        else:
            conv3_kernel_size = 10

        self.conv3 = nn.Sequential(
            nn.Conv2d(
                in_channels=hidden, out_channels=hidden, kernel_size=(1, conv3_kernel_size)
            ),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
            nn.Conv2d(in_channels=hidden, out_channels=hidden, kernel_size=(4, 1)),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
            nn.Conv2d(in_channels=hidden, out_channels=hidden, kernel_size=(4, 1)),
            nn.LeakyReLU(negative_slope=0.01),
            nn.BatchNorm2d(hidden),
        )