diff --git a/dlio_benchmark/data_loader/torch_data_loader.py b/dlio_benchmark/data_loader/torch_data_loader.py index 3048dab1..ffd383de 100644 --- a/dlio_benchmark/data_loader/torch_data_loader.py +++ b/dlio_benchmark/data_loader/torch_data_loader.py @@ -143,7 +143,7 @@ def read(self): batch_size=self.batch_size, sampler=sampler, num_workers=self._args.read_threads, - pin_memory=True, + pin_memory=self._args.pin_memory, drop_last=True, worker_init_fn=dataset.worker_init, **kwargs) @@ -152,7 +152,7 @@ def read(self): batch_size=self.batch_size, sampler=sampler, num_workers=self._args.read_threads, - pin_memory=True, + pin_memory=self._args.pin_memory, drop_last=True, worker_init_fn=dataset.worker_init, **kwargs) # 2 is the default value diff --git a/dlio_benchmark/utils/config.py b/dlio_benchmark/utils/config.py index 13265f41..39ed54d9 100644 --- a/dlio_benchmark/utils/config.py +++ b/dlio_benchmark/utils/config.py @@ -118,6 +118,7 @@ class ConfigArguments: data_loader_sampler: DataLoaderSampler = None reader_classname: str = None multiprocessing_context: str = "fork" + pin_memory: bool = True # derived fields required_samples: int = 1 @@ -521,6 +522,8 @@ def LoadConfig(args, config): args.preprocess_time = reader['preprocess_time'] if 'preprocess_time_stdev' in reader: args.preprocess_time_stdev = reader['preprocess_time_stdev'] + if 'pin_memory' in reader: + args.pin_memory = reader['pin_memory'] # training relevant setting if 'train' in config: diff --git a/docs/source/config.rst b/docs/source/config.rst index 95fecde8..4073509b 100644 --- a/docs/source/config.rst +++ b/docs/source/config.rst @@ -201,6 +201,9 @@ reader * - read_threads* - 1 - number of threads to load the data (for tensorflow and pytorch data loader) + * - pin_memory + - True + - whether to pin the memory for pytorch data loader * - computation_threads - 1 - number of threads to preprocess the data