PrincetonUniversity / PsyNeuLink / 13475398829

85%

master: 85%

Build Type

Pull #3211

github

Committed by web-flow

Commit Message

Add support for batching to AutoDiffCompostion in PyTorch mode. (#3204)

* Changes to allow passing BatchedImpl tensors

When pytorch vmap is used on a function, the function
is called to trace it with tenors passed as inputs
which are not normal pytorch tensors. They are in fact
BatchedImpl tensors that have no underlying storage.
These cannot be converted to double() or numpy arrays
so I needed to remove calls to convert these tensors
in this special case.

* Fix bug in _batch_inputs

_batch_inputs was not actually returning batches of
size batch_size.

* Testing bactch support.

* Fix batch_size=1 case.

* WIP merge for pytorch batching in autodiff

* Fix for when minibatch_size is ndarray

* Cleanup

* Remove handling of batched_results.

* Handle dtype-object in _get_variable_from_input

Not sure this works in general to be honest.

* Fixes from merging in devel.

* Fixes for ragged things.

* More fixes for ragged things.

* Fix bad torch.stack call in execute_input_ports

* Add a real pytorch training test with batching

I made a test that implements an equivalent network
in PyTorch and PNL, initializes them with the same
weights, trains them both and then compares the losses. I think this is similar to some other
tests in the file but this one tests batch size > 1
and it actually trains the pytorch network as well.

* More fixes

* Disable update_autodiff_all_output_values

Now that all_output_values contains the
batch dimension in the first dimension it
is causing and error when passed to
output_CIM.execute because it has the incorrect
number of input ports (unless batch size happens
to match number of input ports). Not sure what
the best fix to this is but disabling this line
seems to work and pass all tests. Need to ask Jon.

* Handle 3D structure better in autodiff update_results.

* Fix formatting

* Fix ragged processing.

A bunch of fixes for things wrong with ragged processing.

* Fix logging test.

We need to add batch dimension to th... (continued)

Pull Request Pull Request #3211: Devel

Run Details

9626 of 12688 branches covered (75.87%)

Branch coverage included in aggregate %.

561 of 630 new or added lines in 23 files covered. (89.05%)

7 existing lines in 5 files now uncovered.

33436 of 38560 relevant lines covered (86.71%)

0.87 hits per line