pytorch-reader: read a checkpoint from bytes, not only from a path
Summary
PytorchReader opens checkpoints by path only. from_reader exists but is documented, and implemented, as plain-pickle only:
The reader must hold a plain pickle: tensor data lives outside the pickle in every PyTorch container, so a checkpoint with tensors must be loaded from a file.
The second half does not follow from the first. Tensor bytes live outside the pickle, inside the container, and a container is just bytes: a ZIP, a legacy stream or a TAR held in a Vec<u8> has everything the file has. Nothing in the parsing needs a path.
Where it bites
- A checkpoint fetched over HTTP, pulled out of an archive, or decrypted, and handed over as bytes, has to be written to a temp file first just to be read.
- The reader's own structure already allows it: a
Containerenum withFile(PathBuf)andMemory(Arc<Vec<u8>>)variants, whoseopen()returns aBox<dyn Read + Seek>over either, is enough for ZIP and legacy, and I have run the parser that way. - The other stores in burn-store (
SafetensorsStore::from_bytes,BurnpackStore::from_bytes) take bytes;PytorchStoreis the odd one out.
Proposed shape
impl PytorchReader {
pub fn from_bytes(bytes: Vec<u8>, top_level_key: Option<&str>) -> Result<Self>;
}Detection already reads only a header; ZipSource can hold ZipArchive<Cursor<Arc<Vec<u8>>>> as easily as ZipArchive<BufReader<File>>, legacy reads seek within the buffer, TAR is already read into memory at open. The Arc matters: every tensor's read closure shares the buffer instead of cloning it. from_reader can then either stay as the plain-pickle entry point with its doc fixed, or read the stream to the end and delegate.
Laziness is unchanged: Tensor::read still gathers on demand, from the buffer instead of the file.
Also
Whatever happens to the constructor, the from_reader doc should say "this function does not read containers" rather than "containers cannot be read from memory".
Source: tracel-ai/burn