diff --git a/active_plugins/runstardist.py b/active_plugins/runstardist.py index 03fe359..b416507 100644 --- a/active_plugins/runstardist.py +++ b/active_plugins/runstardist.py @@ -9,6 +9,7 @@ from stardist.models import StarDist2D, StarDist3D from csbdeep.utils import normalize from numba import cuda +import tensorflow as tf ################################# # @@ -207,6 +208,25 @@ def create_settings(self): maxval=1.0, doc=f"""\ Prevent overlapping +""", + ) + + self.manage_gpu = Binary( + text="Manually set how much GPU memory each worker can use?", + value=False, + doc=""" +If enabled, you can manually set how much of the GPU memory each worker can use. +This is likely to provide the most benefit on Macs. Do not use in a multi-GPU setup.""", + ) + + self.manual_GPU_memory_GB = Float( + text="GPU memory (in GB) for each worker", + value=0.5, + minval=0.0000001, + maxval=30, + doc="""\ +GPU memory in GB available to each worker. Value should be set such that this number times the number +of workers in each copy of CellProfiler times the number of copies of CellProfiler running (if applicable) is <1 """, ) @@ -225,6 +245,8 @@ def settings(self): self.model_choice3D, self.prob_thresh, self.nms_thresh, + self.manage_gpu, + self.manual_GPU_memory_GB, ] def visible_settings(self): @@ -251,7 +273,10 @@ def visible_settings(self): if self.tile_image.value: vis_settings += [self.n_tiles_x, self.n_tiles_y] - vis_settings += [self.prob_thresh, self.nms_thresh, self.gpu_test] + vis_settings += [self.prob_thresh, self.nms_thresh, self.gpu_test, self.manage_gpu] + + if self.manage_gpu.value: + vis_settings += [self.manual_GPU_memory_GB] return vis_settings @@ -272,6 +297,23 @@ def run(self, workspace): raise ValueError( "Greyscale images are not supported by this model. Please provide a color overlay." ) + + # Stolen nearly wholesale from https://wiki.ncsa.illinois.edu/display/ISL20/Managing+GPU+memory+when+using+Tensorflow+and+Pytorch + if self.manage_gpu.value: + # First, Get a list of GPU devices + gpus = tf.config.list_physical_devices('GPU') + if len(gpus) > 0: + # Restrict to only the first GPU. + tf.config.set_visible_devices(gpus[:1], device_type='GPU') + # Create a LogicalDevice with the appropriate memory limit + log_dev_conf = tf.config.LogicalDeviceConfiguration( + memory_limit=self.manual_GPU_memory_GB.value*1024 # 2 GB + ) + # Apply the logical device configuration to the first GPU + tf.config.set_logical_device_configuration( + gpus[0], + [log_dev_conf]) + if self.model.value == CUSTOM_MODEL: model_directory, model_name = os.path.split( @@ -392,11 +434,9 @@ def display(self, workspace, figure): ) def do_check_gpu(self): - import tensorflow - - if len(tensorflow.config.list_physical_devices("GPU")) > 0: + if len(tf.config.list_physical_devices("GPU")) > 0: message = "GPU appears to be working correctly!" - print("GPUs:", tensorflow.config.list_physical_devices("GPU")) + print("GPUs:", tf.config.list_physical_devices("GPU")) else: message = ( "GPU test failed. There may be something wrong with your configuration."