crumpets.torch.shm module

class crumpets.torch.shm.DummyTensorManager(device='cuda:0')[source]

Bases: crumpets.shm.DummyBufferManager

Torch replacement for DummyBufferManager. Returns torch tensors instead of numpy arrays when unpacking.

Parameters

device – output device; buffers are copied here when ready

next()[source]
unpack(data)[source]

Unpack an msgpack message. :param data: msgpack message bytes :return: packed objects

class crumpets.torch.shm.SharedTensorManager(num_buffers, batch_size, buffer_specs, device='cuda:0', _queueclass=<bound method BaseContext.Queue of <multiprocessing.context.DefaultContext object>>)[source]

Bases: crumpets.shm.SharedBufferManager

crumpets.torch.shm.shared_tensor(shape, dtype=<class 'numpy.float32'>, device_type='cuda')[source]

Create a torch tensor that resides in shared memory.

Parameters

  • shape – array shape

  • dtype – numpy dtype

  • device_type – tensor.pin_memory() if ‘cuda’

Returns

np.ndarray