| 6 | |
| 7 | |
| 8 | class augment_inputs_network(nn.Module): |
| 9 | def __init__(self, model): |
| 10 | super(augment_inputs_network, self).__init__() |
| 11 | self.model = model |
| 12 | self.adapter = nn.Conv2d(in_channels=6, out_channels=3, kernel_size=3, stride=1, padding=1) |
| 13 | |
| 14 | def forward(self, x): |
| 15 | x = self.adapter(x) |
| 16 | x = (x - torch.as_tensor(timm.data.constants.IMAGENET_DEFAULT_MEAN, device=x.get_device()).view(1, -1, 1, 1)) / torch.as_tensor(timm.data.constants.IMAGENET_DEFAULT_STD, device=x.get_device()).view(1, -1, 1, 1) |
| 17 | |
| 18 | return self.model(x) |
| 19 | |
| 20 | |
| 21 | class final_model(nn.Module): # Total parameters: 158.64741325378418 MB |
no outgoing calls
no test coverage detected