pytorch-reader: ZIP 向量读取在归档文件周围的 Mutex 上进行序列化; 8 个线程的读取速度不快于 1
作者: antimora创建于 2026年9月17日更新于 2026年9月17日
标签performancestore
use std::os::unix::fs::FileExt; use std::time::Instant; use pytorch_reader::PytorchReader;
fn main() { let path = std::env::args().nth(1).unwrap(); let threads: usize = std::env::args().nth(2).map(|s| s.parse().unwrap()).unwrap_or(8); let reader = PytorchReader::new(&path).unwrap(); let tensors: Vec<_> = reader.tensors().values().cloned().collect(); for t in &tensors { t.read().unwrap(); } // warm the page cache
let best = |f: &dyn Fn()| (0..5).map(|_| { let s = Instant::now(); f(); s.elapsed() }).min().unwrap();
let parallel = best(&|| std::thread::scope(|s| {
for chunk in tensors.chunks(tensors.len().div_ceil(threads)) {
s.spawn(move || for t in chunk { std::hint::black_box(t.read().unwrap()); });
}
}));
println!("Tensor::read, {threads} threads: {parallel:?}");
let file = std::fs::File::open(&path).unwrap();
let mut archive = zip::ZipArchive::new(std::fs::File::open(&path).unwrap()).unwrap();
let ranges: Vec<(u64, usize)> = (0..tensors.len()).map(|i| {
let e = archive.by_name(&format!("archive/data/{i}")).unwrap();
(e.data_start().unwrap(), e.size() as usize)
}).collect();
let pread = best(&|| std::thread::scope(|s| {
for chunk in ranges.chunks(ranges.len().div_ceil(threads)) {
let file = &file;
s.spawn(move || for &(off, len) in chunk {
let mut buf = vec![0u8; len];
file.read_exact_at(off, &mut buf).unwrap();
let mut reader = PytorchReader::new(&path).unwrap();
reader.tensors().values().map(|t| t.read(&mut buf)).collect();
});
}
}));
println!("File::read_exact_at on one shared handle, same byte ranges, no lock: {pread:?}");
}
内容来源: tracel-ai/burn