Skip to content

Commit 0a9c1f5

Browse files
[multiprocessing] Eager tensor support pickle (#48179)
* eager tensot support pickle
1 parent 3f26581 commit 0a9c1f5

File tree

3 files changed

+10
-10
lines changed

3 files changed

+10
-10
lines changed

python/paddle/fluid/dataloader/dataloader_iter.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import itertools
2222
import threading
2323
import numpy as np
24-
import multiprocessing
2524
from collections import namedtuple
2625
from paddle.fluid.framework import (
2726
_set_expected_place,
@@ -422,6 +421,8 @@ def __init__(self, loader):
422421
self._shutdown = False
423422

424423
def _init_workers(self):
424+
import paddle.incubate.multiprocessing as multiprocessing
425+
425426
# multiprocess worker and indice queue list initial as empty
426427
self._workers = []
427428
self._worker_status = []

python/paddle/fluid/dataloader/worker.py

+7-9
Original file line numberDiff line numberDiff line change
@@ -373,21 +373,19 @@ def _worker_loop(
373373
out_queue.put((idx, batch, None))
374374
batch, structure = _flatten_batch(batch)
375375
if use_shared_memory:
376-
# NOTE: In eager mode, Tensor._share_memory has no
377-
# effect, fall back to _array_to_share_memory_tensor
378-
def tensor_share_memory(tensor):
379-
if _in_eager_without_dygraph_check():
380-
return core._array_to_share_memory_tensor(tensor)
381-
return tensor._share_memory()
376+
377+
def numpy2lodtensor(arr):
378+
lodtensor = core.Tensor()
379+
lodtensor.set(arr, core.CPUPlace())
380+
return lodtensor
382381

383382
tensor_list = [
384-
core._array_to_share_memory_tensor(b)
383+
numpy2lodtensor(b)
385384
if isinstance(b, np.ndarray)
386-
else tensor_share_memory(b)
385+
else b.value().get_tensor()
387386
for b in batch
388387
]
389388
out_queue.put((idx, tensor_list, structure))
390-
core._remove_tensor_list_mmap_fds(tensor_list)
391389
else:
392390
out_queue.put((idx, batch, structure))
393391
except KeyboardInterrupt:

python/paddle/incubate/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
from . import autotune # noqa: F401
3535
from . import nn # noqa: F401
3636
from . import asp # noqa: F401
37+
from . import multiprocessing # noqa: F401
3738

3839
from ..fluid.layers.loss import identity_loss
3940

0 commit comments

Comments
 (0)