MIC-DKFZ / medicaldetectiontoolkit

The Medical Detection Toolkit contains 2D + 3D implementations of prevalent object detectors such as Mask R-CNN, Retina Net, Retina U-Net, as well as a training and inference framework focused on dealing with medical images.
Apache License 2.0
1.3k stars 297 forks source link

only a single thread working in "MultiThreadedAugmenter " ! #130

Open wanghan0501 opened 3 years ago

wanghan0501 commented 3 years ago

when I run lidc_exp, I found only a single thread is working while loading data

image

here is my code:

def create_data_gen_pipeline(patient_data, cf, is_training=True):
    """
    create mutli-threaded train/val/test batch generation and augmentation pipeline.
    :param patient_data: dictionary containing one dictionary per patient in the train/test subset.
    :param is_training: (optional) whether to perform data augmentation (training) or not (validation/testing)
    :return: multithreaded_generator
    """

    # create instance of batch generator as first element in pipeline.
    data_gen = BatchGenerator(patient_data, batch_size=cf.batch_size, cf=cf)

    # add transformations to pipeline.
    my_transforms = []
    if is_training:
        mirror_transform = MirrorTransform(axes=np.arange(cf.dim))
        my_transforms.append(mirror_transform)
        spatial_transform = SpatialTransform(patch_size=cf.patch_size[:cf.dim],
                                             patch_center_dist_from_border=cf.da_kwargs['rand_crop_dist'],
                                             do_elastic_deform=cf.da_kwargs['do_elastic_deform'],
                                             alpha=cf.da_kwargs['alpha'], sigma=cf.da_kwargs['sigma'],
                                             do_rotation=cf.da_kwargs['do_rotation'], angle_x=cf.da_kwargs['angle_x'],
                                             angle_y=cf.da_kwargs['angle_y'], angle_z=cf.da_kwargs['angle_z'],
                                             do_scale=cf.da_kwargs['do_scale'], scale=cf.da_kwargs['scale'],
                                             order_data=cf.da_kwargs['order_data'],
                                             random_crop=cf.da_kwargs['random_crop'])

        my_transforms.append(spatial_transform)
    else:
        my_transforms.append(CenterCropTransform(crop_size=cf.patch_size[:cf.dim]))

    my_transforms.append(ConvertSegToBoundingBoxCoordinates(cf.dim, get_rois_from_seg_flag=False, class_specific_seg_flag=cf.class_specific_seg_flag))
    all_transforms = Compose(my_transforms)
    # multithreaded_generator = SingleThreadedAugmenter(data_gen, all_transforms)
    if is_training:
        print (f"is_training: {is_training}, num_processes: {cf.n_workers}")
        # multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes = cf.n_workers, 
        #                          num_cached_per_queue=3, seeds=range(cf.n_workers), pin_memory=True)
        multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes = cf.n_workers, 
        num_cached_per_queue=3, seeds=range(cf.n_workers), pin_memory=True, wait_time=0.01, timeout=5)
    else:
        print (f"is_training: {is_training}, num_processes: {max(1, cf.n_workers // 2)}")
        # multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes = max(1, cf.n_workers // 2),
        #                          num_cached_per_queue=1, seeds=None, pin_memory=True)
        multithreaded_generator = MultiThreadedAugmenter(data_gen, all_transforms, num_processes = max(1, cf.n_workers // 2), 
        num_cached_per_queue=1, seeds=None, pin_memory=True, wait_time=0.01, timeout=5)
    return multithreaded_generator

here is my config:

class configs(DefaultConfigs):

    def __init__(self, server_env=False):

        #########################
        #    Preprocessing      #
        #########################

        # self.root_dir = '/home/gregor/networkdrives/E130-Personal/Goetz/Datenkollektive/Lungendaten/Nodules_LIDC_IDRI'
        # self.raw_data_dir = '{}/new_nrrd'.format(self.root_dir)
        self.pp_dir = '/root/workspace/medicaldetectiontoolkit/datasets/lidc_mdt'
        self.target_spacing = (0.7, 0.7, 1.25)

        #########################
        #         I/O           #
        #########################

        # one out of [2, 3]. dimension the model operates in.
        self.dim = 3

        # one out of ['mrcnn', 'retina_net', 'retina_unet', 'detection_unet', 'ufrcnn'].
        self.model = 'retina_unet'

        DefaultConfigs.__init__(self, self.model, server_env, self.dim)

        #########################
        #         I/O           #
        ######################### 
        self.test_aug = False
        self.n_workers = 4

what's wrong? really thanks for your reply