[dev] support datatype change at split inference for both image and video pipeline

Author: chyomin06

compressai_vision/pipelines/base.py

@@ -81,6 +81,7 @@ def __init__(
         self._create_folder(self.codec_output_dir)
         self.init_time_measure()
         self.init_complexity_measure()
+        self.eval()
 
     def init_time_measure(self):
         self.elapsed_time = {"nn_part_1": 0, "encode": 0, "decode": 0, "nn_part_2": 0}

compressai_vision/pipelines/split_inference/image_split_inference.py

@@ -30,6 +30,7 @@
 import os
 from typing import Dict
 
+import torch
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
@@ -70,6 +71,7 @@ def __init__(
         device: Dict,
     ):
         super().__init__(configs, device)
+        self.datatype = configs["datatype"]
 
     def __call__(
         self,
@@ -116,6 +118,9 @@ def __call__(
                 featureT = self._from_input_to_features(vision_model, d, file_prefix)
                 self.update_time_elapsed("nn_part_1", (time_measure() - start))
 
+                # datatype conversion
+                # featureT["data"] = {k : v.type(getattr(torch, self.datatype)) for k, v in featureT["data"].items()}
+
                 featureT["org_input_size"] = org_img_size
 
                 start = time_measure()

compressai_vision/pipelines/split_inference/video_split_inference.py

@@ -177,10 +177,14 @@ def __call__(
             features["data"] = self._feature_tensor_list_to_dict(
                 self._input_ftensor_buffer
             )
-            # datatype conversion
-            features["data"] = {k : v.type(getattr(torch, self.datatype)) for k, v in features["data"]}
             self._input_ftensor_buffer = []
 
+            # datatype conversion
+            features["data"] = {
+                k: v.type(getattr(torch, self.datatype))
+                for k, v in features["data"].items()
+            }
+
             # Feature Compression
             start = time_measure()
             res, enc_time_by_module, enc_complexity = self._compress(
@@ -250,9 +254,13 @@ def __call__(
 
         # separate a tensor of each keyword item into a list of tensors
         dec_ftensors_list = self._feature_tensor_dict_to_list(dec_features["data"])
-        assert all([self.datatype in str(d.dtype) for d in dec_ftensors_list[0].values()]), "Output features not of expected datatype"
+        assert all(
+            [self.datatype in str(d.dtype) for d in dec_ftensors_list[0].values()]
+        ), "Output features not of expected datatype"
 
-        dec_ftensors_list = [{k : v.type(torch.float32) for k, v in d.items()} for d in dec_ftensors_list]
+        dec_ftensors_list = [
+            {k: v.type(torch.float32) for k, v in d.items()} for d in dec_ftensors_list
+        ]
 
         assert len(dec_ftensors_list) == len(dataloader), (
             f"The number of decoded frames ({len(dec_ftensors_list)}) is not equal "