| 147 | return output |
| 148 | |
| 149 | class CaiTEncoder(ViTModel): |
| 150 | def __init__(self, args, transformer=None, layernorm_epsilon=1e-6, use_final_layernorm=False): |
| 151 | super().__init__(args, transformer=transformer, layernorm_epsilon=layernorm_epsilon, use_final_layernorm=use_final_layernorm) |
| 152 | self.del_mixin('cls') |
| 153 | self.add_mixin('attn', AttnMixin(args.num_attention_heads, args.num_layers)) |
| 154 | self.add_mixin('enc_forward', EncForward(args.hidden_size, args.num_layers, init_values=args.init_scale)) |
| 155 | |
| 156 | @classmethod |
| 157 | def add_model_specific_args(cls, parser): |
| 158 | group = parser.add_argument_group('CaiT-enc', 'CaiT encoder Configurations') |
| 159 | group.add_argument('--init-scale', type=float, default=1e-4) |
| 160 | return super().add_model_specific_args(parser) |
| 161 | |
| 162 | class CaiTDecoder(BaseModel): |
| 163 | def __init__(self, args, transformer=None, layernorm_epsilon=1e-6): |