Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,582 Bytes
476e0f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
from typing import *
from torch import Tensor
import torch
from torch import nn
import torch.nn.functional as tF
from einops import rearrange
from src.options import Options
from src.utils import IMAGENET_MEAN, IMAGENET_STD
class ElevEst(nn.Module):
def __init__(self, opt: Options):
super().__init__()
self.opt = opt
self.backbone: nn.Module = torch.hub.load("facebookresearch/dinov2", opt.elevest_backbone_name)
if opt.freeze_backbone:
self.backbone.requires_grad_(False)
else:
self.backbone.mask_token.requires_grad_(False) # not used
self.dim = dim = {
"dinov2_vits14_reg": 384,
"dinov2_vitb14_reg": 768,
"dinov2_vitl14_reg": 1024,
}[opt.elevest_backbone_name]
self.cls_head = nn.Sequential(
nn.Linear(dim, dim),
nn.GELU(),
nn.Linear(dim, opt.elevest_num_classes),
)
self.offset_head = nn.Sequential(
nn.Linear(dim, dim),
nn.GELU(),
nn.Linear(dim, 1),
)
self.interval = (opt.ele_max - opt.ele_min) / opt.elevest_num_classes
self.register_buffer("lower_bounds", torch.linspace(opt.ele_min, opt.ele_max, opt.elevest_num_classes+1)[:-1])
def state_dict(self, **kwargs):
# Remove frozen parameters without gradients
state_dict = super().state_dict(**kwargs)
if self.opt.freeze_backbone:
for k in list(state_dict.keys()):
if "backbone" in k:
del state_dict[k]
return state_dict
def forward(self, *args, func_name="compute_loss", **kwargs):
# To support different forward functions for models wrapped by `accelerate`
return getattr(self, func_name)(*args, **kwargs)
def compute_loss(self, data: Dict[str, Tensor], dtype: torch.dtype = torch.float32):
outputs = {}
input_images = data["image"].to(dtype) # (B, V, 3, H, W)
gt_elev = data["cam_pose"].to(dtype)[:, :, 0].rad2deg() # (B, V)
input_images = rearrange(input_images, "b v c h w -> (b v) c h w")
gt_elev = rearrange(gt_elev, "b v -> (b v)") # (B*V,)
gt_class = torch.floor((gt_elev - self.opt.ele_min) / self.interval).long()
gt_offset = gt_elev - self.lower_bounds[gt_class]
assert torch.all((gt_class >= 0) & (gt_class < self.opt.elevest_num_classes))
assert torch.all((gt_offset + 1e-8 >= 0) & (gt_offset - 1e-8 < self.interval))
# ImageNet normalization
mean = torch.tensor(IMAGENET_MEAN, device=input_images.device, dtype=dtype).view(3, 1, 1)
std = torch.tensor(IMAGENET_STD, device=input_images.device, dtype=dtype).view(3, 1, 1)
input_images = (input_images - mean) / std
# Predict
features = self.backbone(input_images.to(dtype=dtype), is_training=True)
cls_token = features["x_norm_clstoken"] # (B*V, D)
logits = self.cls_head(cls_token) # (B*V, C)
pred_offset = self.offset_head(cls_token).squeeze(-1).clamp(0., self.interval) # (B*V,)
# Loss
outputs["loss_cls"] = tF.cross_entropy(logits, gt_class)
outputs["loss_offset"] = tF.mse_loss(pred_offset, gt_offset)
outputs["loss"] = outputs["loss_cls"] + self.opt.elevest_reg_weight * outputs["loss_offset"]
with torch.no_grad():
pred_elev = self.lower_bounds[torch.argmax(logits, dim=-1)] + pred_offset # (B*V,)
outputs["err_degree"] = torch.mean(torch.abs(pred_elev - gt_elev))
return outputs
@torch.no_grad()
def predict_elev(self, input_images: Tensor, dtype: torch.dtype = torch.float32):
# Input image preprocessing
input_images = tF.interpolate(input_images, size=(224, 224), mode="bilinear", align_corners=False, antialias=True)
input_images = input_images.to(device=self.lower_bounds.device, dtype=dtype)
mean = torch.tensor(IMAGENET_MEAN, device=input_images.device, dtype=dtype).view(3, 1, 1)
std = torch.tensor(IMAGENET_STD, device=input_images.device, dtype=dtype).view(3, 1, 1)
input_images = (input_images - mean) / std
features = self.backbone(input_images, is_training=True)
cls_token = features["x_norm_clstoken"] # (B, D)
logits = self.cls_head(cls_token) # (B, C)
pred_offset = self.offset_head(cls_token).squeeze(-1).clamp(0., self.interval) # (B,)
pred_elev = self.lower_bounds[torch.argmax(logits, dim=-1)] + pred_offset # (B,)
return pred_elev
|