Pool Resources
Project description
Pool Resources
Description
Library to generalize multiprocessing.Pool(n: int)
to handle generic resources.
First class support for torch devices as a generic resource.
Example for n torch GPUs and k torch Modules
from pool_resources import PoolResources
from pool_resources.resources import TorchResource
class Model(nn.Module):
def __init__(self, input_shape, output_shape):
super().__init__()
self.fc = nn.Linear(input_shape, output_shape)
def forward(self, x):
return self.fc(x)
def forward_fn(item):
model, data = item
return model.forward(data)
modules = [Model(input_shape=20, output_shape=30) for _ in range(k)]
data = tr.randn(k, B, 20)
seq = zip(modules, data)
resources = [TorchResource(f"cuda:{i}") for i in range(n)]
res_sequential = list(map(forward_fn, seq))
res_parallel = PoolResources(resources).map(forward_fn, seq)
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
pool-resources-0.1.tar.gz
(5.8 kB
view hashes)
Built Distribution
Close
Hashes for pool_resources-0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | a0df75b712a716d88ca45b075bffed77a8f5cb92cbc91d6cf9e90b55a01df377 |
|
MD5 | c4fb68d75aaf60d755457cc1e377e583 |
|
BLAKE2b-256 | b5a5858fa1547eca2d2ba60ee7c719205c24c68dc6a567cd9bdc4196ae462c11 |