selitvin
commented
4 years ago Hi.
Second question: the solution is easy: don't have the classes that are supposed to be used on the worker in the main module, as it breaks serialization. I was able to fix it by moving if __name__ == "__main__": clause into a separate module.
In your data structure, you are using only vectors and appears that you do not need multidimensional arrays. In that case, it could be better for you to use vanilla Parquet stores, without Petatorm metadata:
def make_dataset():
def row_generator(x):
return {
"state": [x] * 64,
"action": x % 4,
"reward": 1.,
"mdp_id": x,
"action_probability": [x] * 4,
}
spark = SparkSession.builder.config("spark.driver.memory", "2g").master("local[2]").getOrCreate()
sc = spark.sparkContext
sc.setLogLevel("ERROR")
rows_rdd = (
sc.parallelize(range(DATASET_SIZE))
.map(row_generator)
)
df = spark.createDataFrame(rows_rdd)
df = df.repartition(10)
df.write.mode("overwrite").parquet(PETASTORM_URI)In order to read a vanilla Parquet store with Petastorm, you can:
def time_with_transform():
preprocessor = MockDiscreteDQNBatchPreprocessor()
transform = TransformSpec(transform_row(preprocessor))
data_reader = make_batch_reader(PETASTORM_URI, num_epochs=None, reader_pool_type=READER_POOL_TYPE, transform_spec=transform)
# collate_fn = collate_and_unflatten(PreprocessedTrainingBatch)
with DataLoader(data_reader, batch_size=batch_size) as dataloader:
for _ in range(10):
next(data_reader)
start_time = time.time()
for i, data in enumerate(dataloader):
if i == 0:
print(f"{i} data {type(data)}")
if i >50:
break
print(f"{i} data {type(data)}")
end_time = time.time()
print(f"with transform: {end_time-start_time} secs")- If your consumer is PyTorch code, you don't care about properly declaring your schema. That feature is important for Tensorflow, as we need to know the types to construct computational graph.
- Working with
make_batch_readerwould be faster, as you can do vector processing in transform function on the entire pandas dataframe at once. - Modified time measuring code so it does not include the setup time for the reader and started measuring after some warmup cycles.
- I might have missed your main point about data classes though.... I don't know if this is applicable to your use case, but I think it would be best to leave your data as pd.Dataframe (that is converted to pyarrow table by Petastorm) until you can actually feed it into model, to speed up the system.
Here is how I modified your code.
kaiwenw
Hello, I'm blocked on trying to perform some preprocessing (i.e. transform) on the data before returning from the (Pytorch) dataloader. However, my preprocessing function's return type is a dataclass, rather than a dictionary. My solution is to "flatten" the typed dataclass object to a dictionary in transform, then "unflatten" it back to a dataclass object in collate_fn, which is a bit messy (involves this flattening and unflattening, as well as specifying a new schema of the flattened dictionary). BTW: I also know that I can simply preprocess within collate_fn, but that is only on the master thread and cannot be parallelized; I would prefer to have preprocessing run in parallel.
So first, I was wondering if there's a better way to do this with Petastorm?
In the code example below, I drafted a simple example with this method. When I ran it with reader_pool_type="thread", I saw no parallelization on different cores, which I suspect is from the GIL issue with python. So I tried reader_pool_type="process", which gives me the error "NameError: name 'PreprocessedTrainingBatch' is not defined". Essentially, the global context is not visible in the transform function. For reference, I also tried running examples/mnist/tests/pytorch_example.py with reader_pool_type="process" and it failed with a similar error, saying that the package "transforms" isn't defined.
My second question is if setting reader_pool_type="process" is the way to get it running on multiple cores and if there's a fix to this issue with the transform functions?
Thanks a lot! Kaiwen @selitvin transform.txt