(num_classes: int)
| 361 | |
| 362 | |
| 363 | def cct(num_classes: int): |
| 364 | net = CCT( |
| 365 | img_size=(32, 32), |
| 366 | embedding_dim=256, |
| 367 | n_conv_layers=1, |
| 368 | kernel_size=3, |
| 369 | stride=2, |
| 370 | padding=3, |
| 371 | pooling_kernel_size=3, |
| 372 | pooling_stride=2, |
| 373 | pooling_padding=1, |
| 374 | num_layers=2, |
| 375 | num_heads=2, |
| 376 | mlp_ratio=1., |
| 377 | num_classes=num_classes, |
| 378 | positional_embedding='learnable', # ['sine', 'learnable', 'none'] |
| 379 | ) |
| 380 | return net |