රූපහඳුනාගැනීම සඳහා ගැඹුරු අවශේෂ ඉගෙනීම (රෙස්නෙට්)

57from typing import List, Optional
58
59import torch
60from torch import nn
61
62from labml_helpers.module import Module

65class ShortcutProjection(Module):

72    def __init__(self, in_channels: int, out_channels: int, stride: int):

79        super().__init__()

82        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride)

84        self.bn = nn.BatchNorm2d(out_channels)

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

88        return self.bn(self.conv(x))

91class ResidualBlock(Module):

112    def __init__(self, in_channels: int, out_channels: int, stride: int):

118        super().__init__()

121        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)

123        self.bn1 = nn.BatchNorm2d(out_channels)

125        self.act1 = nn.ReLU()

128        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)

130        self.bn2 = nn.BatchNorm2d(out_channels)

134        if stride != 1 or in_channels != out_channels:

136            self.shortcut = ShortcutProjection(in_channels, out_channels, stride)
137        else:

139            self.shortcut = nn.Identity()

142        self.act2 = nn.ReLU()

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

149        shortcut = self.shortcut(x)

151        x = self.act1(self.bn1(self.conv1(x)))

153        x = self.bn2(self.conv2(x))

155        return self.act2(x + shortcut)

158class BottleneckResidualBlock(Module):

186    def __init__(self, in_channels: int, bottleneck_channels: int, out_channels: int, stride: int):

193        super().__init__()

196        self.conv1 = nn.Conv2d(in_channels, bottleneck_channels, kernel_size=1, stride=1)

198        self.bn1 = nn.BatchNorm2d(bottleneck_channels)

200        self.act1 = nn.ReLU()

203        self.conv2 = nn.Conv2d(bottleneck_channels, bottleneck_channels, kernel_size=3, stride=stride, padding=1)

205        self.bn2 = nn.BatchNorm2d(bottleneck_channels)

207        self.act2 = nn.ReLU()

210        self.conv3 = nn.Conv2d(bottleneck_channels, out_channels, kernel_size=1, stride=1)

212        self.bn3 = nn.BatchNorm2d(out_channels)

216        if stride != 1 or in_channels != out_channels:

218            self.shortcut = ShortcutProjection(in_channels, out_channels, stride)
219        else:

221            self.shortcut = nn.Identity()

224        self.act3 = nn.ReLU()

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

231        shortcut = self.shortcut(x)

233        x = self.act1(self.bn1(self.conv1(x)))

235        x = self.act2(self.bn2(self.conv2(x)))

237        x = self.bn3(self.conv3(x))

239        return self.act3(x + shortcut)

242class ResNetBase(Module):

256    def __init__(self, n_blocks: List[int], n_channels: List[int],
257                 bottlenecks: Optional[List[int]] = None,
258                 img_channels: int = 3, first_kernel_size: int = 7):

267        super().__init__()

270        assert len(n_blocks) == len(n_channels)

273        assert bottlenecks is None or len(bottlenecks) == len(n_channels)

277        self.conv = nn.Conv2d(img_channels, n_channels[0],
278                              kernel_size=first_kernel_size, stride=2, padding=first_kernel_size // 2)

280        self.bn = nn.BatchNorm2d(n_channels[0])

283        blocks = []

285        prev_channels = n_channels[0]

287        for i, channels in enumerate(n_channels):

290            stride = 2 if len(blocks) == 0 else 1
291
292            if bottlenecks is None:

294                blocks.append(ResidualBlock(prev_channels, channels, stride=stride))
295            else:

298                blocks.append(BottleneckResidualBlock(prev_channels, bottlenecks[i], channels,
299                                                      stride=stride))

302            prev_channels = channels

304            for _ in range(n_blocks[i] - 1):
305                if bottlenecks is None:

307                    blocks.append(ResidualBlock(channels, channels, stride=1))
308                else:

310                    blocks.append(BottleneckResidualBlock(channels, bottlenecks[i], channels, stride=1))

313        self.blocks = nn.Sequential(*blocks)

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

321        x = self.bn(self.conv(x))

323        x = self.blocks(x)

325        x = x.view(x.shape[0], x.shape[1], -1)

327        return x.mean(dim=-1)