gpu_cache

args.py

@@ -24,7 +24,8 @@ parser.add_argument('--skip_model', action='store_true')
 parser.add_argument('--ifm', type=str)
 parser.add_argument('--anime4k', action='store_true')
 parser.add_argument('--alpha_split', action='store_true')
-parser.add_argument('--cache', type=str, default='100mb')
+parser.add_argument('--cache', type=str, default='256mb')
+parser.add_argument('--gpu_cache', type=str, default='512mb')
 parser.add_argument('--simplify', type=int, default=1)
 args = parser.parse_args()
 args.output_w = int(args.output_size.split('x')[0])
@@ -33,4 +34,8 @@ if args.cache is not None:
     args.max_cache_len=int(convert_to_byte(args.cache)/262144)
 else:
     args.max_cache_len=0
+if args.gpu_cache is not None:
+    args.max_gpu_cache_len=int(convert_to_byte(args.gpu_cache)/589824)
+else:
+    args.max_gpu_cache_len=0
 if args.output_webcam is None and args.output_dir is None: args.debug = True

main.py

@@ -165,6 +165,7 @@ class ModelClientProcess(Process):
         self.model_fps_number = Value('f', 0.0)
         self.gpu_fps_number = Value('f', 0.0)
         self.cache_hit_ratio = Value('f', 0.0)
+        self.gpu_cache_hit_ratio= Value('f', 0.0)
 
     def run(self):
         model = None
@@ -302,22 +303,23 @@ class ModelClientProcess(Process):
             input_hash = hash(tuple(model_input))
             cached = model_cache.get(input_hash)
             tot += 1
+            mouth_eye_vector_c=[0.0]*27
             if not cached is None:
                 self.output_queue.put_nowait(cached)
                 model_cache.move_to_end(input_hash)
                 hit += 1
-                self.cache_hit_ratio.value = hit / tot
             else:
                 if args.perf:
                     tic = time.perf_counter()
                 for i in range(27):
                     mouth_eye_vector[0, i] = model_input[i]
+                    mouth_eye_vector_c[i]=model_input[i]
                 for i in range(3):
                     pose_vector[0, i] = model_input[i + 27]
                 if model is None:
                     output_image = input_image
                 else:
-                    output_image = model(input_image, mouth_eye_vector, pose_vector)
+                    output_image = model(input_image, mouth_eye_vector, pose_vector, mouth_eye_vector_c,self.gpu_cache_hit_ratio)
                 if args.perf:
                     torch.cuda.synchronize()
                     print("model", (time.perf_counter() - tic) * 1000)
@@ -337,9 +339,10 @@ class ModelClientProcess(Process):
                 if args.max_cache_len > 0:
                     model_cache[input_hash] = postprocessed_image
                     if len(model_cache) > args.max_cache_len:
-                        model_cache.popitem(0)
+                        model_cache.popitem(last=False)
             if args.debug:
                 self.model_fps_number.value = model_fps()
+                self.cache_hit_ratio.value = hit / tot
 
 
 @torch.no_grad()
@@ -639,7 +642,10 @@ def main():
                 cv2.putText(output_frame, str('IFM_FPS:%.1f' % client_process.ifm_fps_number.value), (0, 48),
                             cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
             if args.max_cache_len > 0:
-                cv2.putText(output_frame, str('CACHED:%.1f%%' % (model_process.cache_hit_ratio.value * 100)), (0, 64),
+                cv2.putText(output_frame, str('MEMCACHED:%.1f%%' % (model_process.cache_hit_ratio.value * 100)), (0, 64),
+                            cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
+            if args.max_gpu_cache_len > 0:
+                cv2.putText(output_frame, str('GPUCACHED:%.1f%%' % (model_process.gpu_cache_hit_ratio.value * 100)), (0, 80),
                             cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
             cv2.imshow("frame", output_frame)
             # cv2.imshow("camera", debug_image)

models.py

@@ -6,18 +6,36 @@ from tha2.poser.modes.mode_20 import load_face_morpher, load_face_rotater, load_
 
 from args import args
 
+from collections import OrderedDict
+
 class TalkingAnimeLight(nn.Module):
     def __init__(self):
         super(TalkingAnimeLight, self).__init__()
         self.face_morpher = load_face_morpher('pretrained/face_morpher.pt')
         self.two_algo_face_rotator = load_face_rotater('pretrained/two_algo_face_rotator.pt')
         self.combiner = load_combiner('pretrained/combiner.pt')
+        self.face_cache = OrderedDict()
+        self.tot = 0
+        self.hit = 0
 
-    def forward(self, image, mouth_eye_vector, pose_vector):
+    def forward(self, image, mouth_eye_vector, pose_vector, mouth_eye_vector_c, ratio=None):
         x = image.clone()
         if args.perf:
             tic=time.perf_counter()
-        mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector)
+        input_hash = hash(tuple(mouth_eye_vector_c))
+        cached = self.face_cache.get(input_hash)
+        self.tot+=1
+        if cached is None:
+            mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector)
+            self.face_cache[input_hash]=mouth_eye_morp_image.detach()
+            if len(self.face_cache) > args.max_gpu_cache_len:
+                self.face_cache.popitem(last=False)
+        else:
+            self.hit+=1
+            mouth_eye_morp_image=cached
+            self.face_cache.move_to_end(input_hash)
+        if args.debug and ratio is not None:
+            ratio.value=self.hit / self.tot
         if args.perf:
             print(" - face_morpher",(time.perf_counter()-tic)*1000)
             tic=time.perf_counter()