From 825a28244c03b534f868936640b0b0ef722e3ca1 Mon Sep 17 00:00:00 2001
From: LDOUBLEV <liuvv0203@outlook.com>
Date: Mon, 11 Jan 2021 15:57:28 +0800
Subject: [PATCH] add use_shared_memory params for AIstudio demo

---
 configs/det/det_mv3_db.yml        | 4 +++-
 configs/rec/rec_icdar15_train.yml | 2 ++
 ppocr/data/__init__.py            | 7 ++++---
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/configs/det/det_mv3_db.yml b/configs/det/det_mv3_db.yml
index bdb4afc0..00a16b5a 100644
--- a/configs/det/det_mv3_db.yml
+++ b/configs/det/det_mv3_db.yml
@@ -102,6 +102,7 @@ Train:
     drop_last: False
     batch_size_per_card: 16
     num_workers: 8
+    use_shared_memory: False
 
 Eval:
   dataset:
@@ -128,4 +129,5 @@ Eval:
     shuffle: False
     drop_last: False
     batch_size_per_card: 1 # must be 1
-    num_workers: 8
\ No newline at end of file
+    num_workers: 8
+    use_shared_memory: False
\ No newline at end of file
diff --git a/configs/rec/rec_icdar15_train.yml b/configs/rec/rec_icdar15_train.yml
index 3de0ce77..8a743b50 100644
--- a/configs/rec/rec_icdar15_train.yml
+++ b/configs/rec/rec_icdar15_train.yml
@@ -76,6 +76,7 @@ Train:
     batch_size_per_card: 256
     drop_last: True
     num_workers: 8
+    use_shared_memory: False
 
 Eval:
   dataset:
@@ -96,3 +97,4 @@ Eval:
     drop_last: False
     batch_size_per_card: 256
     num_workers: 4
+    use_shared_memory: False
diff --git a/ppocr/data/__init__.py b/ppocr/data/__init__.py
index 7b0faf12..4e1ff0ae 100644
--- a/ppocr/data/__init__.py
+++ b/ppocr/data/__init__.py
@@ -66,8 +66,10 @@ def build_dataloader(config, mode, device, logger):
     batch_size = loader_config['batch_size_per_card']
     drop_last = loader_config['drop_last']
     num_workers = loader_config['num_workers']
-
-    use_shared_memory = False
+    if 'use_shared_memory' in loader_config.keys():
+        use_shared_memory = loader_config['use_shared_memory']
+    else:
+        use_shared_memory = True
     if mode == "Train":
         #Distribute data to multiple cards
         batch_sampler = DistributedBatchSampler(
@@ -75,7 +77,6 @@ def build_dataloader(config, mode, device, logger):
             batch_size=batch_size,
             shuffle=False,
             drop_last=drop_last)
-        use_shared_memory = True
     else:
         #Distribute data to single card
         batch_sampler = BatchSampler(
-- 
GitLab