Efficient Torch Tensor Serialization

[shikokuchuo]

Inviting @dfalbel

Continuing our email conversation here so it is easier to keep track.

Current Status Update:

nanonext 0.10.0 (released on CRAN) completes the optimisation work for data serialization in general.

• R objects are directly serialised into NNG's send buffer or unserialized from NNG's receive buffer
• Eliminates intermediate data copies
• Made possible by using serialisation at R's C interface

New Feature: Allow efficient serialisation / unserialization of Torch Tensors

• Uses Safetensors serialization (or can be switched to whatever the 'official' method is e.g. from Torch package)
• Uses R serialization's built-in 'refhook' system to record references of tensor objects - allows perfect recreation of objects without the need for additional marshalling
• Serialised tensors are appended to the raw vector using the custom serialization implementation in nanonext (not dissimilar to how Safetensors work)
• All using optimal serialisation methods - no slowdown due to ASCII serialisation etc.
• Enabled at the underlying binding level (here at nanonext) - meaning it 'just works' in mirai without requiring any change in that package

Initial testing implementation: https://github.com/shikokuchuo/nanonext/tree/tensor

I've updated the examples here: https://github.com/shikokuchuo/wip/blob/main/tensor.R

To dos:

• Current branch is a working demo, but interface will change - to either fully flexible to allow for all kinds of serialisation (not just Torch Tensors), or specialise and optimize heavily
• Quite a few call backs into R from C - inherited from R's refhook design which allows user-supplied functions. Could be eliminated for performance (subject to the above API design point)
• Safetensors at the moment only offers an R interface. If this becomes the bottleneck then it should be straightforward to re-factor to interface with nanonext at the C level. I have some ideas here which I can share if this becomes relevant.

Next Steps:

To discuss how best to utilise mirai in the Torch / mlverse ecosystem.

[shikokuchuo]

Copying @lionel- I mentioned at the Sprint. Also mirai.promises is a CRAN package with repo here: https://github.com/shikokuchuo/mirai.promises

[wlandau]

Sounds exciting.

NB if you plan to expose torch (un)serialization methods directly to R users, I might adopt them for the built-in torch storage format class in targets: https://github.com/ropensci/targets/blob/main/R/class_torch.R. You will probably outperform my existing S3 methods store_marshal_object.tar_torch() and store_unmarshal_object.tar_torch().

[shikokuchuo]

If exposing directly to users, this would belong in a new package. And I think this would need to have built-in the most common formats rather than expecting users (as opposed to developers) to supply their own. What other formats apart from Torch are important for Targets? The name for this new package currently eludes me though...

[wlandau]

If exposing directly to users, this would belong in a new package.

If that's the case, then please don't worry about it. Sounds like a lot more work than I expected, and on reflection it would create back-compatibility challenges for targets because of how existing torch data is currently saved.

And I think this would need to have built-in the most common formats rather than expecting users (as opposed to developers) to supply their own. What other formats apart from Torch are important for Targets? The name for this new package currently eludes me though...

targets has built-in storage formats for these cases. They are mostly general, covering external files and data frames, and the "keras" and "torch" formats are a bit specialized (almost too specialized). (If I were to write targets again, these would be left out of the package and left to users to implement using tar_format()).

[shikokuchuo]

I see, yes I don't think it's a particular fit for saving targets. Where it would come into its own is in marshalling / unmarshalling. If the relevant 'refhook' functions are registered then the user would be able to send his entire workspace to a remote worker, even it contained nested objects such as dataframes of torch tensors (?!)...

[shikokuchuo]

A note to @dfalbel and others - this functionality was merged in ad5af25 and is no longer experimental.

The next CRAN release of {nanonext} allows send / recv of reference objects such as Torch tensors across R sessions.

In the case of {torch}, this would be as simple as the following one-time call to register the serialization/unserialization functions:

nextmode(safetensors::safe_serialize, safetensors::safe_load_file)

This capability will make its way into {mirai} in due course.

[shikokuchuo]

This functionality has been implemented in next_config() in the nanonext 0.11.0 release.

Also implemented is the mirai integration. A rather rudimentary example is included in the package vignette here:
https://shikokuchuo.net/mirai/articles/mirai.html#serialization-custom-functions
It shows the sending/receiving of torch tensors across processes.

@dfalbel I've added a line to credit our discussions in the package readme - I assume you're ok with that.

The next CRAN release of mirai is imminent. After that it will probably be best to continue discussions in that forum, so I will close the issue here.

[shikokuchuo]

Closed with the release of mirai 0.11.3. Please see the readme https://shikokuchuo.net/mirai/index.html#use-with-torch or the vignette https://shikokuchuo.net/mirai/articles/mirai.html#serialization-custom-functions