| 47 | return out |
| 48 | |
| 49 | class YOLOS(ViTModel): |
| 50 | def __init__(self, args, transformer=None, layernorm_epsilon=1e-6, **kwargs): |
| 51 | super().__init__(args, transformer=transformer, layernorm_epsilon=layernorm_epsilon, **kwargs) |
| 52 | self.del_mixin('patch_embedding') |
| 53 | self.add_mixin('patch_embedding', NewTokenMixin(args.vocab_size+args.num_det_tokens, args.in_channels, args.hidden_size, self.property)) |
| 54 | self.del_mixin('cls') |
| 55 | self.add_mixin('det_head', DetHeadMixin(args)) |
| 56 | |
| 57 | @classmethod |
| 58 | def add_model_specific_args(cls, parser): |
| 59 | group = parser.add_argument_group('YOLOS', 'YOLOS Configurations') |
| 60 | group.add_argument('--num-det-tokens', type=int) |
| 61 | group.add_argument('--num-det-classes', type=int) |
| 62 | return super().add_model_specific_args(parser) |