[torch] Retain a Tensor object past enclosing PointerScope?

[guangster]

Hi, I would like to persist a Tensor object outside a PoinsterScope. Following your example here
gives me access to Tensor but the underlying data seems to get cleaned. The only solution I found is to call retainReference() on the original Pointer I passed to torch.from_blob(Pointer,...). This is not ideal as I have to track each Pointer I used to create the Tensor. Neither calling retainReference() on Tensor, nor on the output of Tensor.data_ptr_float() worked (corrupted data after) even though Tensor.data_ptr_float() returned a pointer with the same address as the original Pointer. Is there a way for us to get the original Pointer from the tensor, or some better ways to retain the tensor?

I attached my example below along with the output I observed. Any help is greatly appreciated!

import org.bytedeco.javacpp.FloatPointer;
import org.bytedeco.javacpp.PointerScope;
import org.bytedeco.pytorch.Tensor;
import org.bytedeco.pytorch.global.torch;

import java.util.Arrays;

public class PointerScopeTest {
    public static void main(String[] args) {
        Tensor tensor;
        FloatPointer pointer;
        float[] buf = new float[3];

        try (PointerScope scope = new PointerScope()) {

            FloatPointer fp = new FloatPointer(new float[]{3, 2, 1});
            tensor = torch.from_blob(fp);
            pointer = new FloatPointer(new float[]{3, 2, 1});
            tensor.retainReference();
            pointer.retainReference();
            // fp.retainReference(); // I would like to not have to call this
            // because I would like to only keep track of tensor.

            
            // work with tensor here...
            // ...


            // sanity check, data still look ok here
            pointer.get(buf);
            System.out.println(Arrays.toString(buf)); //[3.0, 2.0, 1.0]
            tensor.data_ptr_float().get(buf);
            System.out.println(Arrays.toString(buf)); //[3.0, 2.0, 1.0]
        }

        // I would like to be able to access tensor here.
        // It seems the objects are available.
        System.out.println(tensor.isNull()); // false
        System.out.println(pointer.isNull()); // false

        // for Pointer, the data is still ok:
        pointer.get(buf);
        System.out.println(Arrays.toString(buf)); //[3.0, 2.0, 1.0]

        // for Tensor however, the data appears corrupted if I comment out "fp.retainReference();" above
        tensor.data_ptr_float().get(buf);
        System.out.println(Arrays.toString(buf)); //[8.464593E-5, 4.5852E-41, 7.964554E-5]


    }
}

[saudet]

You could copy the data to a new Tensor instead. You could also transfer the ownership of the Pointer to the new Tensor by providing a deallocator to from_blob().

[guangster]

Thanks! Your new commit worked with Loader.addressof("free"). I can now call retainReference on both the tensor and the pointer, then just close the tensor to close the data as well. Out of curiosity, how would we pass directly values like Pointer.NativeDeallocator.deallocatorAddress since it is a private field? Would we have to use reflection to pull the field out, or is there a better way?

Finally, I attached how I am using your commit below, let me know if I missed anything.
Thanks again!

  public static void main(String[] args) {
        Tensor tensor1, tensor2;
        FloatPointer hidden1, hidden2;
        float[] buf = new float[3];
        float[] data = new float[]{3,2,1};

        try (PointerScope scope = new PointerScope()) {

            hidden1 = new FloatPointer(data);
            hidden2 = new FloatPointer(data);

            tensor1 = torch.for_blob(hidden1).deleter(Loader.addressof("free"))
                    .make_tensor();
            tensor2 = torch.for_blob(hidden2).make_tensor();

            tensor1.retainReference();
            tensor2.retainReference();
            hidden1.retainReference();
            hidden2.retainReference();

        }
        hidden1.get(buf);
        System.out.println("tensor1's data before cleanup:" + Arrays.toString(buf));

        hidden2.get(buf);
        System.out.println("tensor2's data before cleanup:" + Arrays.toString(buf));

        tensor1.close();
        tensor2.close();
        
        hidden1.get(buf);
        System.out.println("tensor1's data after cleanup:" + Arrays.toString(buf));

        hidden2.get(buf);
        System.out.println("tensor2's data after cleanup:" + Arrays.toString(buf));
        
        //tensor1's data before cleanup:[3.0, 2.0, 1.0]
        //tensor2's data before cleanup:[3.0, 2.0, 1.0]
        //tensor1's data after cleanup:[0.0, 0.0, -1.0844938E-19]
        //tensor2's data after cleanup:[3.0, 2.0, 1.0]

    }

[saudet]

Right, it wasn't intended to be used that way, but I don't see how adding getters could harm anything, so I've done that in commit bytedeco/javacpp@730ae11. With that, it can be accessed this way now:

    class MyFloatPointer extends FloatPointer {
        long deallocatorAddress() {
            return ((NativeDeallocator)deallocator()).deallocatorAddress();
        }
    }

If you want to keep using free() though, that's only guaranteed to work with memory allocated by malloc() and friends, so in that case you should use Pointer.malloc() to allocate memory, as done in the samples for Triton Inference Server: https://github.com/bytedeco/javacpp-presets/tree/master/tritonserver/samples

BTW, call deallocate(false) instead of retainReference(). With the latter, the deallocator may still get called by the GC.

[guangster]

You are right, I noticed my Tensor objects were getting garbage collected when I expose them outside Java (i.e. to Python for bytedeco/javacpp#530), but they worked okay in Java. I ended up setting org.bytedeco.javacpp.nopointergc to true and that solved it.

I took a look at deallocate(false) and logically that makes sense what it would do, but in practice, this didn't quite work for torch's tensor object, as this makes the data_ptr object inaccessible (NullPointerException: This pointer address is NULL). The Tensor is still fine if I call deallocate(false) on hidden and retainReference() on tensor and I can operate with it. I can even call item_float() to look at the first element which looks okay, but I cannot get the result out to Java by accessing its data_ptr_float(). Only way I found to get the data out is if I call torch.clone(tensor2).data_ptr_float(), which shows me that it is working as intended (data is okay outside scope) but not really suitable for my usecase.
Let me know if I have missed something, but if not, it is still not a problem for me because using nopointergc works well for me.

[saudet]

You need to call deallocate(false) on the original FloatPointer object that has a deallocator. It's not going to work if you call it on the Tensor or on Tensor.data_ptr().

[saudet]

It's actually also necessary to call retainReference() since it doesn't remove the deallocator, but only disables it for the GC, so close() can still call it.