RGBD-SOD/dptdepth

1from typing import Dict, Optional2 3from torch import Tensor, nn4from transformers import PreTrainedModel5 6from .configuration_dptdepth import DPTDepthConfig7from .models import DPTDepthModel as DPTDepth8 9 10class DPTDepthModel(PreTrainedModel):11    """12    The line that sets the config_class is not mandatory,13    unless you want to register your model with the auto classes14    """15 16    config_class = DPTDepthConfig17 18    def __init__(self, config: DPTDepthConfig):19        super().__init__(config)20        self.model = DPTDepth()21        self.loss = nn.L1Loss()22 23    """24    You can have your model return anything you want, 25    but returning a dictionary with the loss included when labels are passed, 26    will make your model directly usable inside the Trainer class. 27    Using another output format is fine as long as you are planning on 28    using your own training loop or another library for training.29    """30 31    def forward(self, rgbs: Tensor, gts: Optional[Tensor] = None) -> Dict[str, Tensor]:32        logits = self.model(rgbs)33        if gts is not None:34            loss = self.loss(logits, gts)35            return {"loss": loss, "logits": logits}36        return {"logits": logits}37