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