| 36 | |
| 37 | |
| 38 | class DPTHead(nn.Module): |
| 39 | def __init__( |
| 40 | self, |
| 41 | in_channels, |
| 42 | features=256, |
| 43 | use_bn=False, |
| 44 | out_channels=[256, 512, 1024, 1024], |
| 45 | use_clstoken=False |
| 46 | ): |
| 47 | super(DPTHead, self).__init__() |
| 48 | |
| 49 | self.use_clstoken = use_clstoken |
| 50 | |
| 51 | self.projects = nn.ModuleList([ |
| 52 | nn.Conv2d( |
| 53 | in_channels=in_channels, |
| 54 | out_channels=out_channel, |
| 55 | kernel_size=1, |
| 56 | stride=1, |
| 57 | padding=0, |
| 58 | ) for out_channel in out_channels |
| 59 | ]) |
| 60 | |
| 61 | self.resize_layers = nn.ModuleList([ |
| 62 | nn.ConvTranspose2d( |
| 63 | in_channels=out_channels[0], |
| 64 | out_channels=out_channels[0], |
| 65 | kernel_size=4, |
| 66 | stride=4, |
| 67 | padding=0), |
| 68 | nn.ConvTranspose2d( |
| 69 | in_channels=out_channels[1], |
| 70 | out_channels=out_channels[1], |
| 71 | kernel_size=2, |
| 72 | stride=2, |
| 73 | padding=0), |
| 74 | nn.Identity(), |
| 75 | nn.Conv2d( |
| 76 | in_channels=out_channels[3], |
| 77 | out_channels=out_channels[3], |
| 78 | kernel_size=3, |
| 79 | stride=2, |
| 80 | padding=1) |
| 81 | ]) |
| 82 | |
| 83 | if use_clstoken: |
| 84 | self.readout_projects = nn.ModuleList() |
| 85 | for _ in range(len(self.projects)): |
| 86 | self.readout_projects.append( |
| 87 | nn.Sequential( |
| 88 | nn.Linear(2 * in_channels, in_channels), |
| 89 | nn.GELU())) |
| 90 | |
| 91 | self.scratch = _make_scratch( |
| 92 | out_channels, |
| 93 | features, |
| 94 | groups=1, |
| 95 | expand=False, |