| |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from .base import BaseModel |
| from .feature_extractor import FeatureExtractor |
|
|
|
|
| class MapEncoder(BaseModel): |
| default_conf = { |
| "embedding_dim": "???", |
| "output_dim": None, |
| "num_classes": "???", |
| "backbone": "???", |
| "unary_prior": False, |
| } |
|
|
| def _init(self, conf): |
| self.embeddings = torch.nn.ModuleDict( |
| { |
| k: torch.nn.Embedding(n + 1, conf.embedding_dim) |
| for k, n in conf.num_classes.items() |
| } |
| ) |
| |
| input_dim = len(conf.num_classes) * conf.embedding_dim |
| output_dim = conf.output_dim |
| if output_dim is None: |
| output_dim = conf.backbone.output_dim |
| if conf.unary_prior: |
| output_dim += 1 |
| if conf.backbone is None: |
| self.encoder = nn.Conv2d(input_dim, output_dim, 1) |
| elif conf.backbone == "simple": |
| self.encoder = nn.Sequential( |
| nn.Conv2d(input_dim, 128, 3, padding=1), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(128, 128, 3, padding=1), |
| nn.ReLU(inplace=True), |
| nn.Conv2d(128, output_dim, 3, padding=1), |
| ) |
| else: |
| self.encoder = FeatureExtractor( |
| { |
| **conf.backbone, |
| "input_dim": input_dim, |
| "output_dim": output_dim, |
| } |
| ) |
|
|
| def _forward(self, data): |
| embeddings = [ |
| self.embeddings[k](data["map"][:, i]) |
| for i, k in enumerate(("areas", "ways", "nodes")) |
| ] |
| embeddings = torch.cat(embeddings, dim=-1).permute(0, 3, 1, 2) |
| if isinstance(self.encoder, BaseModel): |
| features = self.encoder({"image": embeddings})["feature_maps"] |
| else: |
| features = [self.encoder(embeddings)] |
| pred = {} |
| if self.conf.unary_prior: |
| pred["log_prior"] = [f[:, -1] for f in features] |
| features = [f[:, :-1] for f in features] |
| pred["map_features"] = features |
| return pred |
|
|
