Method forward

model/model_utils.py:397–431

Perform dynamic predictions.

(self, feats, timestamp=None)

    
    @autocast('cuda', enabled=False)
    def forward(self, feats, timestamp=None):
        """
        Perform dynamic predictions.
        """
        feats = feats.type(torch.float32)
        feats = rearrange(feats, 'b v n d -> (b v) n d')
        gsparams = {}
        prior_params = {}
        for key in ["xyz_dynamic", "opacity_dynamic"]:
            v = feats
            if f'{key}_scale' in self.gs_layer and key == "xyz_dynamic":
                logits_pred = torch.ones(feats.shape[0], feats.shape[1], 1).to(feats.device).float()
                means = self.gs_layer[key](v) 
                log_scales = self.gs_layer[f'{key}_scale'](v)
                logits, means, log_scales = self.prior.expand_params(logits_pred, means, log_scales, mean_activation='tanh')
                prior_params[key] = {"logits": logits, "means": means, "log_scales": log_scales}  # [B, N*r, dim, nr_mix]
                val, probs = self.prior.sample(logits, means, log_scales)
                val = val.reshape(*val.shape[:2], -1, 3)  # [B, N, L * forder, 3]
                val = val * self.dynamic_scalar
                gsparams[key] = val
            elif f'{key}_scale' in self.gs_layer and key == "opacity_dynamic":
                logits_pred = torch.ones(feats.shape[0], feats.shape[1], 1).to(feats.device).float()
                val = self.gs_layer[key](v) 
                log_scales = self.gs_layer[f'{key}_scale'](v)
                scalar = torch.exp(val[..., 0:1])
                means = val[..., 1:2]
                logits, means, log_scales = self.prior.expand_params(logits_pred, means, log_scales, mean_activation='tanh')
                prior_params[key] = {"logits": logits, "means": means, "log_scales": log_scales}  # [B, N*r, dim, nr_mix]
                val, probs = self.prior.sample(logits, means, log_scales)
                val = 0.5 + 0.5 * val  # t1 in [0, 1]
                gsparams[key] = torch.cat([scalar, val], dim=-1)
            else:
                v = self.gs_layer[key](v)
                gsparams[key] = self.key_activation(v, key)
        return gsparams, prior_params

    @autocast('cuda', enabled=False)
    def key_activation(self, v: torch.Tensor, key=''):