Merge pull request #36 from Tobias-Fischer/pytorch_environment

Pytorch environment

rt_gene/README.md

@@ -23,8 +23,8 @@ More information can be found on the Personal Robotic Lab's website: <https://ww
     - For kinetic: `sudo apt-get install python-catkin-tools ros-kinetic-ros-numpy ros-kinetic-camera-info-manager-py ros-kinetic-uvc-camera`
     - For melodic: `sudo apt-get install python-catkin-tools ros-melodic-uvc-camera`
 1. Install required Python packages:
-    - For `conda` users (recommended): `conda install tensorflow-gpu keras numpy scipy tqdm pytorch torchvision pillow && conda install -c 1adrianb face_alignment && pip install empy rospkg psutil`
-    - For `pip` users: `pip install tensorflow-gpu keras numpy scipy tqdm torch torchvision Pillow face-alignment psutil`
+    - For `conda` users (recommended): `conda install tensorflow-gpu keras numpy scipy tqdm pytorch torchvision pillow && conda install -c 1adrianb face_alignment && pip install empy rospkg`
+    - For `pip` users: `pip install tensorflow-gpu keras numpy scipy tqdm torch torchvision Pillow face-alignment`
 1. Download and build RT-GENE:
     1. `cd $HOME/catkin_ws/src && git clone https://github.com/Tobias-Fischer/rt_gene.git`
     1. `cd $HOME/catkin_ws && catkin build`

rt_gene/launch/estimate_gaze.launch

@@ -1,4 +1,5 @@
 <launch>
+    <arg name="pytorch_num_threads" default="8" />
     <arg name="use_face_encoding_tracker" default="True" />
     <arg name="interpupillary_distance" default="0.058"/>
     <arg name="start_rviz" default="False"/>
@@ -17,6 +18,10 @@
     <arg name="device_id_facealignment"  default="cuda:0" /> <!-- pyTorch format, e.g. cuda:0 or cpu:0 -->
     <arg name="device_id_gazeestimation"  default="/gpu:0" /> <!-- tensorflow format, e.g. /gpu:0 or /cpu:0 -->
 
+    <env name="OMP_NUM_THREADS" value="$(arg pytorch_num_threads)" />
+    <env name="KMP_SETTING" value="KMP_AFFINITY=granularity=fine,compact,1,0" />
+    <env name="KMP_BLOCKTIME" value="1" />
+
     <group ns="$(arg group_name)">
       <node pkg="tf2_ros" type="static_transform_publisher" name="kinect2_link_broadcaster"
             args="0.0 0.0 0.0 -0.5 0.5 -0.5 0.5 $(arg rgb_frame_id_ros) $(arg rgb_frame_id)" if="$(arg is_frame_rotated)" />

rt_gene/scripts/estimate_gaze.py

@@ -17,7 +17,6 @@
 from tqdm import tqdm
 
 import rospy
-import psutil
 from sensor_msgs.msg import Image
 from cv_bridge import CvBridge
 import rospkg
@@ -46,10 +45,7 @@ class GazeEstimator(object):
     Additionally, two images with the gaze overlaid on the eye images are published."""
 
     def __init__(self):
-        os.environ["OMP_NUM_THREADS"] = str(psutil.cpu_count(logical=False))
-        os.environ["KMP_SETTING"] = "KMP_AFFINITY=granularity=fine,compact,1,0"
-        os.environ["KMP_BLOCKTIME"] = "1"
-
+        tqdm.write("PyTorch using {} threads.".format(os.environ["OMP_NUM_THREADS"]))
         self.image_height = rospy.get_param("~image_height", 36)
         self.image_width = rospy.get_param("~image_width", 60)
         self.bridge = CvBridge()