projects/deep_video_compression/_utils.py [135:174]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        return flow, image2_est, probabilities

    def compute_batch_loss(self, image1, image2):
        assert self.loss_functions.entropy_fn is not None
        assert image1.ndim == image2.ndim == 4

        flow, image2_est, probabilities = self.forward(image1, image2)

        # compute distortion loss
        distortion_loss = self.loss_functions.distortion_fn(image2, image2_est)

        # compute compression loss, average over num pixels
        num_pixels = image1.shape[0] * image1.shape[-2] * image1.shape[-1]
        entropy_loss = self.loss_functions.entropy_fn(probabilities, num_pixels)

        return (
            LossValues(
                distortion_loss=distortion_loss,
                flow_entropy_loss=entropy_loss,
                resid_entropy_loss=None,
            ),
            OutputTensors(
                flow=flow.detach(),
                image1=image1.detach(),
                image2=image2.detach(),
                image2_est=image2_est.detach(),
            ),
        )

    def quantile_loss(self):
        return self.motion_entropy_bottleneck.loss()

    def update(self, force=False):
        return self.motion_entropy_bottleneck.update(force=force)

    def recompose_model(self, model):
        model.motion_estimator = self.motion_estimator
        model.motion_encoder = self.motion_encoder
        model.motion_entropy_bottleneck = self.motion_entropy_bottleneck
        model.motion_decoder = self.motion_decoder
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



projects/deep_video_compression/_utils.py [206:245]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        return flow, image2_est, probabilities

    def compute_batch_loss(self, image1, image2):
        assert self.loss_functions.entropy_fn is not None
        assert image1.ndim == image2.ndim == 4

        flow, image2_est, probabilities = self.forward(image1, image2)

        # compute distortion loss
        distortion_loss = self.loss_functions.distortion_fn(image2, image2_est)

        # compute compression loss, average over num pixels
        num_pixels = image1.shape[0] * image1.shape[-2] * image1.shape[-1]
        entropy_loss = self.loss_functions.entropy_fn(probabilities, num_pixels)

        return (
            LossValues(
                distortion_loss=distortion_loss,
                flow_entropy_loss=entropy_loss,
                resid_entropy_loss=None,
            ),
            OutputTensors(
                flow=flow.detach(),
                image1=image1.detach(),
                image2=image2.detach(),
                image2_est=image2_est.detach(),
            ),
        )

    def quantile_loss(self):
        return self.motion_entropy_bottleneck.loss()

    def update(self, force=False):
        return self.motion_entropy_bottleneck.update(force=force)

    def recompose_model(self, model):
        model.motion_estimator = self.motion_estimator
        model.motion_encoder = self.motion_encoder
        model.motion_entropy_bottleneck = self.motion_entropy_bottleneck
        model.motion_decoder = self.motion_decoder
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -