tanshuai0219 / EDTalk

[ECCV 2024 Oral] EDTalk - Official PyTorch Implementation
Apache License 2.0
309 stars 26 forks source link

How to optimize CPU usage by modifying test_EDTalk_A_use_exp_weight function to leverage GPU #31

Open emboss33 opened 2 weeks ago

emboss33 commented 2 weeks ago

I noticed a significant increase in CPU usage in the following part of demo_EDTalk_A_using_predefined_exp_weights.py

I suspect that the function test_EDTalk_A_use_exp_weight in generator.py is causing this spike in CPU usage. I am planning to modify this function to utilize GPU (CUDA) more efficiently in order to reduce the load on the CPU and improve the calculation speed, anyone here?

img_target_lip = self.lip_vid_target[i:i+1]
if i >= len_pose:
    img_target_pose = self.pose_vid_target[:, -1, :, :, :]
else:
    img_target_pose = self.pose_vid_target[:, i, :, :, :]

img_recon = self.gen.test_EDTalk_A_use_exp_weight(self.img_source, img_target_lip, img_target_pose, self.exp_vid_target, h_start)

vid_target_recon.append(img_recon.unsqueeze(2))
tanshuai0219 commented 2 weeks ago

I noticed a significant increase in CPU usage in the following part of demo_EDTalk_A_using_predefined_exp_weights.py

I suspect that the function test_EDTalk_A_use_exp_weight in generator.py is causing this spike in CPU usage. I am planning to modify this function to utilize GPU (CUDA) more efficiently in order to reduce the load on the CPU and improve the calculation speed, anyone here?

img_target_lip = self.lip_vid_target[i:i+1]
if i >= len_pose:
    img_target_pose = self.pose_vid_target[:, -1, :, :, :]
else:
    img_target_pose = self.pose_vid_target[:, i, :, :, :]

img_recon = self.gen.test_EDTalk_A_use_exp_weight(self.img_source, img_target_lip, img_target_pose, self.exp_vid_target, h_start)

vid_target_recon.append(img_recon.unsqueeze(2))

I have checked the codes, and the current models and function calculation are on GPU. I think only loading data depends on CPU.