Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[pytorch] Add NDList to IValue unit test #1762

Merged
merged 1 commit into from
Jul 1, 2022
Merged

[pytorch] Add NDList to IValue unit test #1762

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 59 additions & 8 deletions engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,26 +12,27 @@
*/
package ai.djl.pytorch.jni;

import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.Shape;
import ai.djl.pytorch.engine.PtNDArray;
import ai.djl.pytorch.engine.PtNDManager;

import org.testng.Assert;
import org.testng.annotations.Test;

import java.util.Arrays;

public class IValueUtilsTest {

@Test
public void getInputsTestTupleSyntax() {
try (PtNDManager manager = (PtNDManager) NDManager.newBaseManager()) {
PtNDArray array1 = (PtNDArray) manager.zeros(new Shape(1));
array1.setName("Test()");
PtNDArray array2 = (PtNDArray) manager.ones(new Shape(1));
array2.setName("Test()");
public void testTuple() {
try (NDManager manager = NDManager.newBaseManager()) {
NDArray array1 = manager.zeros(new Shape(1));
array1.setName("input1()");
NDArray array2 = manager.ones(new Shape(1));
array2.setName("input1()");
NDList input = new NDList(array1, array2);
// the NDList is mapped to (input1: Tuple(Tensor))
input.attach(manager);

IValue[] iValues = IValueUtils.getInputs(input);
Expand All @@ -42,4 +43,54 @@ public void getInputsTestTupleSyntax() {
Arrays.stream(iValues).forEach(IValue::close);
}
}

@Test
public void testMapOfTensor() {
try (NDManager manager = NDManager.newBaseManager()) {
NDArray array1 = manager.zeros(new Shape(1));
array1.setName("input1.key1");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: add some comments to explain why use input1 as prefix? Like group all keys with same prefix into one Map?

NDArray array2 = manager.ones(new Shape(1));
array2.setName("input1.key2");
NDArray array3 = manager.zeros(new Shape(1));
array3.setName("input2.key1");
NDArray array4 = manager.ones(new Shape(1));
array4.setName("input2.key2");
NDArray array5 = manager.ones(new Shape(1));
array5.setName("input2.key3");
NDList input = new NDList(array1, array2, array3, array4, array5);
// the NDList is mapped to (input1: Dict(str, Tensor), input2: Dict(str, Tensor))
Copy link
Contributor

@ylwu-amzn ylwu-amzn Jul 2, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any way to generate something like

(input1: Dict(str, list(Tensor)), input2: Dict(str, list(Tensor)))

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. We don't support this currently, it's possible to add nested IValue support.
  2. But you have to use IValue directly, and use PyTorch specific API, which doesn't work with Translator pattern. something like: https://github.com/deepjavalibrary/djl/blob/master/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java#L184-L191

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I created a PR for your issue: #1765

// the first part of NDArray name is the variable name of the inputs
// the 2nd part is the key of each value in the dict
input.attach(manager);

IValue[] iValues = IValueUtils.getInputs(input);
Assert.assertEquals(iValues.length, 2);
Assert.assertTrue(iValues[0].isMap());
Assert.assertEquals(iValues[1].toIValueMap().size(), 3);

Arrays.stream(iValues).forEach(IValue::close);
}
}

@Test
public void testListOfTensor() {
try (NDManager manager = NDManager.newBaseManager()) {
NDArray array1 = manager.zeros(new Shape(1));
array1.setName("input1[]");
NDArray array2 = manager.ones(new Shape(1));
array2.setName("input1[]");
NDArray array3 = manager.ones(new Shape(1));
array3.setName("input2");
NDList input = new NDList(array1, array2, array3);
// the NDList is mapped to (input1: list(Tensor), input2: Tensor)
input.attach(manager);

IValue[] iValues = IValueUtils.getInputs(input);
Assert.assertEquals(iValues.length, 2);
Assert.assertTrue(iValues[0].isList());
Assert.assertEquals(iValues[0].toIValueArray().length, 2);

Arrays.stream(iValues).forEach(IValue::close);
}
}
}