yjguo:

import os
import time
import argparse
from torch.multiprocessing import Process, Queue, Lock
import torch

parser = argparse.ArgumentParser(description='data loader')
parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
                    help='number of workers (default: 2)')

def get_ten(tens, lock):
    #lock.acquire()
    if not tens.empty():
        ten = tens.get()
    else:
        ten = None
    #lock.release()
    return ten

def put_ten(tens, ten, lock):
    #lock.acquire()
    tens.put(ten)
    #lock.release()

def worker(tens, process_index, tenlock):
    #affinity_mask = {process_index+1}
    #os.sched_setaffinity(0, affinity_mask)
    #torch.set_num_threads(1)
    while True:
        time.sleep(0.005)
        ten = torch.ones([720, 1280, 3])
        put_ten(tens, ten, tenlock)

def main():
    args = parser.parse_args()
    #affinity_mask = {0}
    #os.sched_setaffinity(0, affinity_mask)
    #torch.set_num_threads(1)

    tens = Queue()
    process_count = args.workers
    process_list = []
    tenlock = Lock()
    for i in range(process_count):
        p = Process(target=worker, args=(tens, i, tenlock))
        p.start()
        process_list.append(p)

    ten_count = 0
    while True:
        ten = get_ten(tens, tenlock)
        if ten is not None:
            ten_count = ten_count + 1
            if ten_count % 500 == 0:
                print("in consumer process, get %d tensors, still %d left in the queue" % (ten_count, tens.qsize()))
        else:
            print("in consumer process, do not get a tensor")
            time.sleep(0.01)

    for p in process_list:
        p.join()

if __name__ == '__main__':
    main()



When I uncomment these lines:
    #affinity_mask = {process_index+1}
    #os.sched_setaffinity(0, affinity_mask)
    #torch.set_num_threads(1)

I can get the following outputs w/o OOM. I think not setting CPU affinity caused the slowdown you’ve seen.
in consumer process, do not get a tensor
in consumer process, do not get a tensor
in consumer process, do not get a tensor
in consumer process, get 500 tensors, still 222 left in the queue
in consumer process, get 1000 tensors, still 223 left in the queue
in consumer process, get 1500 tensors, still 223 left in the queue
in consumer process, get 2000 tensors, still 219 left in the queue
in consumer process, get 2500 tensors, still 223 left in the queue
in consumer process, get 3000 tensors, still 223 left in the queue
in consumer process, get 3500 tensors, still 219 left in the queue
in consumer process, get 4000 tensors, still 224 left in the queue
in consumer process, get 4500 tensors, still 225 left in the queue
in consumer process, get 5000 tensors, still 231 left in the queue