| 314 | # CCT Main model |
| 315 | |
| 316 | class CCT(nn.Module): |
| 317 | def __init__( |
| 318 | self, |
| 319 | img_size=224, |
| 320 | embedding_dim=768, |
| 321 | n_input_channels=3, |
| 322 | n_conv_layers=1, |
| 323 | kernel_size=7, |
| 324 | stride=2, |
| 325 | padding=3, |
| 326 | pooling_kernel_size=3, |
| 327 | pooling_stride=2, |
| 328 | pooling_padding=1, |
| 329 | *args, **kwargs |
| 330 | ): |
| 331 | super().__init__() |
| 332 | img_height, img_width = pair(img_size) |
| 333 | |
| 334 | self.tokenizer = Tokenizer(n_input_channels=n_input_channels, |
| 335 | n_output_channels=embedding_dim, |
| 336 | kernel_size=kernel_size, |
| 337 | stride=stride, |
| 338 | padding=padding, |
| 339 | pooling_kernel_size=pooling_kernel_size, |
| 340 | pooling_stride=pooling_stride, |
| 341 | pooling_padding=pooling_padding, |
| 342 | max_pool=True, |
| 343 | activation=nn.ReLU, |
| 344 | n_conv_layers=n_conv_layers, |
| 345 | conv_bias=False) |
| 346 | |
| 347 | self.classifier = TransformerClassifier( |
| 348 | sequence_length=self.tokenizer.sequence_length(n_channels=n_input_channels, |
| 349 | height=img_height, |
| 350 | width=img_width), |
| 351 | embedding_dim=embedding_dim, |
| 352 | seq_pool=True, |
| 353 | dropout_rate=0., |
| 354 | attention_dropout=0.1, |
| 355 | stochastic_depth=0.1, |
| 356 | *args, **kwargs) |
| 357 | |
| 358 | def forward(self, x): |
| 359 | x = self.tokenizer(x) |
| 360 | return self.classifier(x) |
| 361 | |
| 362 | |
| 363 | def cct(num_classes: int): |