(self, croconet)
| 34 | self.layer_dims = layer_dims |
| 35 | |
| 36 | def setup(self, croconet): |
| 37 | dpt_args = {'output_width_ratio': self.output_width_ratio, 'num_channels': self.num_channels} |
| 38 | if self.hooks_idx is None: |
| 39 | if hasattr(croconet, 'dec_blocks'): # encoder + decoder |
| 40 | step = {8: 3, 12: 4, 24: 8}[croconet.dec_depth] |
| 41 | hooks_idx = [croconet.dec_depth+croconet.enc_depth-1-i*step for i in range(3,-1,-1)] |
| 42 | else: # encoder only |
| 43 | step = croconet.enc_depth//4 |
| 44 | hooks_idx = [croconet.enc_depth-1-i*step for i in range(3,-1,-1)] |
| 45 | self.hooks_idx = hooks_idx |
| 46 | print(f' PixelwiseTaskWithDPT: automatically setting hook_idxs={self.hooks_idx}') |
| 47 | dpt_args['hooks'] = self.hooks_idx |
| 48 | dpt_args['layer_dims'] = self.layer_dims |
| 49 | self.dpt = DPTOutputAdapter(**dpt_args) |
| 50 | dim_tokens = [croconet.enc_embed_dim if hook<croconet.enc_depth else croconet.dec_embed_dim for hook in self.hooks_idx] |
| 51 | dpt_init_args = {'dim_tokens_enc': dim_tokens} |
| 52 | self.dpt.init(**dpt_init_args) |
| 53 | |
| 54 | |
| 55 | def forward(self, x, img_info): |
no test coverage detected