| 210 | |
| 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] |
| 231 | x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] |
| 232 | x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] |
| 233 | x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width] |
| 234 | x = x + self.positional_embedding.to(x.dtype) |
| 235 | x = self.ln_pre(x) |
| 236 | |
| 237 | x = x.permute(1, 0, 2) # NLD -> LND |
| 238 | x = self.transformer(x) |
| 239 | x = x.permute(1, 0, 2) # LND -> NLD |
| 240 | |
| 241 | x = self.ln_post(x[:, 0, :]) |
| 242 | |
| 243 | if self.proj is not None: |
| 244 | x = x @ self.proj |
| 245 | |
| 246 | return x |
| 247 | |
| 248 | |
| 249 | class CLIP(nn.Module): |