| | import torch |
| | from torch import nn |
| | from torch.nn import functional as F |
| |
|
| |
|
| | class ChannelLastConv1d(nn.Conv1d): |
| |
|
| | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| | x = x.permute(0, 2, 1) |
| | x = super().forward(x) |
| | x = x.permute(0, 2, 1) |
| | return x |
| |
|
| |
|
| | |
| | class MLP(nn.Module): |
| |
|
| | def __init__( |
| | self, |
| | dim: int, |
| | hidden_dim: int, |
| | multiple_of: int = 256, |
| | ): |
| | """ |
| | Initialize the FeedForward module. |
| | |
| | Args: |
| | dim (int): Input dimension. |
| | hidden_dim (int): Hidden dimension of the feedforward layer. |
| | multiple_of (int): Value to ensure hidden dimension is a multiple of this value. |
| | |
| | Attributes: |
| | w1 (ColumnParallelLinear): Linear transformation for the first layer. |
| | w2 (RowParallelLinear): Linear transformation for the second layer. |
| | w3 (ColumnParallelLinear): Linear transformation for the third layer. |
| | |
| | """ |
| | super().__init__() |
| | hidden_dim = int(2 * hidden_dim / 3) |
| | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| |
|
| | self.w1 = nn.Linear(dim, hidden_dim, bias=False) |
| | self.w2 = nn.Linear(hidden_dim, dim, bias=False) |
| | self.w3 = nn.Linear(dim, hidden_dim, bias=False) |
| |
|
| | def forward(self, x): |
| | return self.w2(F.silu(self.w1(x)) * self.w3(x)) |
| |
|
| |
|
| | class ConvMLP(nn.Module): |
| |
|
| | def __init__( |
| | self, |
| | dim: int, |
| | hidden_dim: int, |
| | multiple_of: int = 256, |
| | kernel_size: int = 3, |
| | padding: int = 1, |
| | ): |
| | """ |
| | Initialize the FeedForward module. |
| | |
| | Args: |
| | dim (int): Input dimension. |
| | hidden_dim (int): Hidden dimension of the feedforward layer. |
| | multiple_of (int): Value to ensure hidden dimension is a multiple of this value. |
| | |
| | Attributes: |
| | w1 (ColumnParallelLinear): Linear transformation for the first layer. |
| | w2 (RowParallelLinear): Linear transformation for the second layer. |
| | w3 (ColumnParallelLinear): Linear transformation for the third layer. |
| | |
| | """ |
| | super().__init__() |
| | hidden_dim = int(2 * hidden_dim / 3) |
| | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| |
|
| | self.w1 = ChannelLastConv1d(dim, |
| | hidden_dim, |
| | bias=False, |
| | kernel_size=kernel_size, |
| | padding=padding) |
| | self.w2 = ChannelLastConv1d(hidden_dim, |
| | dim, |
| | bias=False, |
| | kernel_size=kernel_size, |
| | padding=padding) |
| | self.w3 = ChannelLastConv1d(dim, |
| | hidden_dim, |
| | bias=False, |
| | kernel_size=kernel_size, |
| | padding=padding) |
| |
|
| | def forward(self, x): |
| | return self.w2(F.silu(self.w1(x)) * self.w3(x)) |
| |
|
