පැච්ඔබට අවශ්ය සියලුම ද? (කොන්වී මික්සර්)

38import torch
39from torch import nn
40
41from labml_helpers.module import Module
42from labml_nn.utils import clone_module_list

45class ConvMixerLayer(Module):

54    def __init__(self, d_model: int, kernel_size: int):

59        super().__init__()

63        self.depth_wise_conv = nn.Conv2d(d_model, d_model,
64                                         kernel_size=kernel_size,
65                                         groups=d_model,
66                                         padding=(kernel_size - 1) // 2)

68        self.act1 = nn.GELU()

70        self.norm1 = nn.BatchNorm2d(d_model)

74        self.point_wise_conv = nn.Conv2d(d_model, d_model, kernel_size=1)

76        self.act2 = nn.GELU()

78        self.norm2 = nn.BatchNorm2d(d_model)

80    def forward(self, x: torch.Tensor):

82        residual = x

85        x = self.depth_wise_conv(x)
86        x = self.act1(x)
87        x = self.norm1(x)

90        x += residual

93        x = self.point_wise_conv(x)
94        x = self.act2(x)
95        x = self.norm2(x)

98        return x

101class PatchEmbeddings(Module):

110    def __init__(self, d_model: int, patch_size: int, in_channels: int):

116        super().__init__()

121        self.conv = nn.Conv2d(in_channels, d_model, kernel_size=patch_size, stride=patch_size)

123        self.act = nn.GELU()

125        self.norm = nn.BatchNorm2d(d_model)

127    def forward(self, x: torch.Tensor):

132        x = self.conv(x)

134        x = self.act(x)
135        x = self.norm(x)

138        return x

141class ClassificationHead(Module):

151    def __init__(self, d_model: int, n_classes: int):

156        super().__init__()

158        self.pool = nn.AdaptiveAvgPool2d((1, 1))

160        self.linear = nn.Linear(d_model, n_classes)

162    def forward(self, x: torch.Tensor):

164        x = self.pool(x)

166        x = x[:, :, 0, 0]

168        x = self.linear(x)

171        return x

174class ConvMixer(Module):

181    def __init__(self, conv_mixer_layer: ConvMixerLayer, n_layers: int,
182                 patch_emb: PatchEmbeddings,
183                 classification: ClassificationHead):

191        super().__init__()

193        self.patch_emb = patch_emb

195        self.classification = classification

197        self.conv_mixer_layers = clone_module_list(conv_mixer_layer, n_layers)

199    def forward(self, x: torch.Tensor):

204        x = self.patch_emb(x)

207        for layer in self.conv_mixer_layers:
208            x = layer(x)

211        x = self.classification(x)

214        return x