facebookresearch / playtorch

PlayTorch is a framework for rapidly creating mobile AI experiences.
https://playtorch.dev/
MIT License
828 stars 102 forks source link

Tensor assign at index #172

Closed lusinlu closed 1 year ago

lusinlu commented 1 year ago

Version

No response

Problem Area

react-native-pytorch-core (core package)

Steps to Reproduce

Hi! if Im not wrong, currently it is not possible to assign a tensor at the index, i.e. this code will fail on the last line of assignment

images = // some array of Images
batch_size = 10  
var batch_tensor = torch.zeros([batch_size, 3, 224, 224])
for(var idx=0; idx<images.length;idx++)
    image = images[idx]
    height = image.getHeight();
    width = image.getWidth();

    const blob = media.toBlob(image);
    let tensor = torch.fromBlob(blob, [height, width, 3]);
    tensor = tensor.permute([2, 0, 1]);
    tensor = resize(tensor);
    batch_tensor[idx] = tensor

Expected Results

Is it possible to have a tensor assign operation to work, or are there some other methods for creating a batch of images?

Code example, screenshot, or link to repository

No response

raedle commented 1 year ago

@lusinlu, that's correct. It is not possible to assign a tensor using the Tensor Indexing API.

However, it should be possible by implementing the Tensor.index_put function in TensorHostObject by overriding the JSI HostObject.set function.

This is analogous to the Tensor.index implementation in the TensorHostObject.get function.

I can advise if you want to take a stab and submit a PR!

raedle commented 1 year ago

@lusinlu, alternatively, it is possible to concatenate tensors with the torch.cat function.

const tensor1 = torch.tensor(
  [
    [
      [1, 2],
      [3, 4],
      [5, 6],
    ],
  ],
  { dtype: torch.uint8 }
);
const tensor2 = torch.tensor(
  [
    [
      [7, 8],
      [9, 10],
      [11, 12],
    ],
  ],
  { dtype: torch.uint8 }
);
const tensor = torch.cat([tensor1, tensor2]);
console.log(tensor.shape);
console.log([...tensor.data()]);
[2,3,2]
[1,2,3,4,5,6,7,8,9,10,11,12]

Check this PlayTorch example: https://snack.playtorch.dev/@raedle/tensor_cat

raedle commented 1 year ago

I submitted PR #174 as working draft for the "inplace put" Tensor Indexing API

lusinlu commented 1 year ago

Hi @raedle ! Many thanks for such a prompt response and for the implementation. Indeed, the concat also can do the job, thanks for pointing out to it.