| 1423 | |
| 1424 | |
| 1425 | def _add_vit_args(parser): |
| 1426 | group = parser.add_argument_group(title="vit") |
| 1427 | |
| 1428 | group.add_argument( |
| 1429 | "--num-classes", |
| 1430 | type=int, |
| 1431 | default=1000, |
| 1432 | help="num of classes in vision classificaiton task", |
| 1433 | ) |
| 1434 | group.add_argument( |
| 1435 | "--img-dim", |
| 1436 | type=int, |
| 1437 | default=224, |
| 1438 | help="Image size for vision classification task", |
| 1439 | ) |
| 1440 | group.add_argument( |
| 1441 | "--num-channels", |
| 1442 | type=int, |
| 1443 | default=3, |
| 1444 | help="Number of channels in input image data", |
| 1445 | ) |
| 1446 | group.add_argument( |
| 1447 | "--patch-dim", type=int, default=16, help="patch dimension used in vit" |
| 1448 | ) |
| 1449 | |
| 1450 | return parser |
| 1451 | |
| 1452 | |
| 1453 | def _add_zero_args(parser): |