(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int)
| 205 | |
| 206 | class VisionTransformer(nn.Module): |
| 207 | def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int): |
| 208 | super().__init__() |
| 209 | self.input_resolution = input_resolution |
| 210 | self.output_dim = output_dim |
| 211 | self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False) |
| 212 | |
| 213 | scale = width ** -0.5 |
| 214 | self.class_embedding = nn.Parameter(scale * torch.randn(width)) |
| 215 | self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width)) |
| 216 | self.ln_pre = LayerNorm(width) |
| 217 | |
| 218 | self.transformer = Transformer(width, layers, heads) |
| 219 | |
| 220 | self.ln_post = LayerNorm(width) |
| 221 | self.proj = nn.Parameter(scale * torch.randn(width, output_dim)) |
| 222 | |
| 223 | def forward(self, x: torch.Tensor): |
| 224 | x = self.conv1(x) # shape = [*, width, grid, grid] |
nothing calls this directly
no test coverage detected