Spaces:
Sleeping
Sleeping
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| import torch.nn as nn | |
| from torchvision.models.resnet import Bottleneck | |
| from .base import BaseModel | |
| from .feature_extractor import AdaptationBlock | |
| from .utils import checkpointed | |
| class BEVNet(BaseModel): | |
| default_conf = { | |
| "pretrained": True, | |
| "num_blocks": "???", | |
| "latent_dim": "???", | |
| "input_dim": "${.latent_dim}", | |
| "output_dim": "${.latent_dim}", | |
| "confidence": False, | |
| "norm_layer": "nn.BatchNorm2d", # normalization ind decoder blocks | |
| "checkpointed": False, # whether to use gradient checkpointing | |
| "padding": "zeros", | |
| } | |
| def _init(self, conf): | |
| blocks = [] | |
| Block = checkpointed(Bottleneck, do=conf.checkpointed) | |
| for i in range(conf.num_blocks): | |
| dim = conf.input_dim if i == 0 else conf.latent_dim | |
| blocks.append( | |
| Block( | |
| dim, | |
| conf.latent_dim // Bottleneck.expansion, | |
| norm_layer=eval(conf.norm_layer), | |
| ) | |
| ) | |
| self.blocks = nn.Sequential(*blocks) | |
| self.output_layer = AdaptationBlock(conf.latent_dim, conf.output_dim) | |
| if conf.confidence: | |
| self.confidence_layer = AdaptationBlock(conf.latent_dim, 1) | |
| def update_padding(module): | |
| if isinstance(module, nn.Conv2d): | |
| module.padding_mode = conf.padding | |
| if conf.padding != "zeros": | |
| self.bocks.apply(update_padding) | |
| def _forward(self, data): | |
| features = self.blocks(data["input"]) | |
| pred = { | |
| "output": self.output_layer(features), | |
| } | |
| if self.conf.confidence: | |
| pred["confidence"] = self.confidence_layer(features).squeeze(1).sigmoid() | |
| return pred | |
| def loss(self, pred, data): | |
| raise NotImplementedError | |
| def metrics(self, pred, data): | |
| raise NotImplementedError | |