(self, opt: Options, **model_kwargs)
| 449 | |
| 450 | class GSPredictor(nn.Module): |
| 451 | def __init__(self, opt: Options, **model_kwargs): |
| 452 | super(GSPredictor, self).__init__() |
| 453 | self.opt = opt |
| 454 | self.encoder = GSEncoder(opt, **model_kwargs) |
| 455 | self.patch_size = opt.patch_size |
| 456 | |
| 457 | if len(opt.down_resolution) > 0: |
| 458 | self.actual_input_res = opt.down_resolution |
| 459 | else: |
| 460 | self.actual_input_res = (opt.image_height, opt.image_width) |
| 461 | |
| 462 | # Upsampler |
| 463 | self.decoder_ratio = opt.decoder_ratio |
| 464 | if self.decoder_ratio > 0: |
| 465 | self.gaussian_upsampler = GaussianUpsampler(width=opt.hidden_dim, |
| 466 | ch_decay=2, |
| 467 | up_ratio=(2**self.decoder_ratio), |
| 468 | low_channels=opt.decoder_dim, |
| 469 | window_size=opt.window_size, |
| 470 | opt=opt) |
| 471 | transformer_dim = self.gaussian_upsampler.out_channels |
| 472 | else: |
| 473 | transformer_dim = opt.hidden_dim |
| 474 | # check if exist gaussian upsampler |
| 475 | # if hasattr(self.encoder, "gaussian_upsampler"): |
| 476 | # transformer_dim = self.encoder.gaussian_upsampler.out_channels |
| 477 | # else: |
| 478 | # transformer_dim = self.opt.hidden_dim |
| 479 | if self.opt.use_pm: |
| 480 | self.predictor = GSPMDecoder(opt, transformer_dim=transformer_dim) |
| 481 | else: |
| 482 | raise NotImplementedError("Currently only support probabilistic predictor") |
| 483 | |
| 484 | def _freeze(self): |
| 485 | for name, param in self.named_parameters(): |
nothing calls this directly
no test coverage detected