Two independent linear maps using Conv1d

dual_linear.py

"Two independent linear maps using grouped Conv1d"
import torch
import torch.nn as nn


class DualLinearSerial(nn.Module):
    def __init__(self, in_channels, out_channels, bias=True):
        super().__init__()
        self.left = nn.Linear(in_channels // 2, out_channels // 2, bias=bias)
        self.right = nn.Linear(in_channels // 2, out_channels // 2, bias=bias)

    def forward(self, x):
        l, r = torch.tensor_split(x, 2, dim=-1) # split channels in half
        return torch.cat([self.left(l), self.right(r)], dim=-1)


class DualLinearParallel(nn.Module):
    def __init__(self, in_channels, out_channels, bias=True):
        super().__init__()
        self.lr = nn.Conv1d(in_channels, out_channels, kernel_size=1, groups=2, bias=bias)

    def forward(self, x):
        return self.lr(x)


if __name__ == '__main__':
    x = torch.randn(1, 3, 6) # N T C

    m1 = DualLinearSerial(6, 14)
    y1 = m1.forward(x) # Linear assumes N T C
    print('m1', y1.shape) # (1, 3, 14)

    m2 = DualLinearParallel(6, 14)

    print('m1 weights', m1.left.weight.shape, m1.left.bias.shape) # (7, 1) (7,)
    print('m2 weights', m2.lr.weight.shape, m2.lr.bias.shape) # (14, 1, 1) (14,)

    m2.lr.weight = nn.Parameter(torch.cat([m1.left.weight, m1.right.weight], dim=0).unsqueeze(-1))
    m2.lr.bias = nn.Parameter(torch.cat([m1.left.bias, m1.right.bias], dim=0))
    
    y2 = m2.forward(x.mT).mT # Conv1d assumes N C T
    print('m2', y2.shape) # (1, 3, 14)
    assert torch.allclose(y1, y2)