GitLab

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/557

To avoid circular dependencies, move the function `add_distillation_configs` that defines the default config for a `runner` making use of distillation from `mobile-vision/d2go/d2go/modeling/distillation.py` to `mobile-vision/d2go/d2go/runner/config_defaults.py`.

Reviewed By: wat3rBro

Differential Revision: D46096374

fbshipit-source-id: eb85d91b5239e7ab10809a9bf84c869d05d32401

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/455

The test can be flaky due to numerical mismatch if using `self.AssertEqual`, eg. https://www.internalfb.com/intern/testinfra/diagnostics/1688850007977704.562950031998292.1672749571/

```
Traceback (most recent call last):
  File "/data/sandcastle/boxes/eden-trunk-hg-fbcode-fbsource/buck-out/v2/gen/fbcode/104a4d5c3a690252/mobile-vision/d2go/tests/__modeling_test_modeling_distillation__/modeling_test_modeling_distillation#link-tree/d2go/tests/modeling/test_modeling_distillation.py", line 674, in test_da_train
    self.assertEqual(
AssertionError: {'rea[14 chars]2894], grad_fn=<MulBackward0>), 'synthetic': t[85 chars]d0>)} != {'rea[14 chars]2894]), 'synthetic': tensor([1.4532]), 'add': [13 chars]64])}
- {'add': tensor([18.0064], grad_fn=<MulBackward0>),
-  'real': tensor([0.2894], grad_fn=<MulBackward0>),
-  'synthetic': tensor([1.4532], grad_fn=<MulBackward0>)}
+ {'add': tensor([18.0064]),
+  'real': tensor([0.2894]),
+  'synthetic': tensor([1.4532])}
```

.Change to use `torch.testing.assert_close` instead for tensor comparison.

Reviewed By: YanjunChen329

Differential Revision: D42352509

fbshipit-source-id: 8a647685d1347a9bd493f2faed7e066eb9159e14

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/439

As title

Reviewed By: mattcyu1

Differential Revision: D41759804

fbshipit-source-id: 929efa960be570f0fe8543600e012d1bf037ab3b

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/432

We support caching of tuples since they behave similarly to lists

Reviewed By: XiaoliangDai

Differential Revision: D41483876

fbshipit-source-id: 9d741074f8e2335ddd737ae3f1bdb288910f5564

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/431

Add a generic domain adaptation algorithm. This algorithm:
* gets domain0 data out of the dataloader
* runs domain0 data into the model and saves target layer output
* gets domain1 data of the dataloader
* runs domain1 data into the model and saves target layer output
* runs domain adaptation loss on domain0, domain1 outputs
* combines losses using model training iteration

This diffs adds `get_preprocess_domain0_input` and `get_preprocess_domain1_input` to the distillation helper. These are functions that the user can use to convert the dataloader output to something that will be used by the model (e.g., pull the domain0 or domain1 key out of a dataloader that returns a dict).

Differential Revision: D40970724

fbshipit-source-id: fff050fbe864654fa6cb0df927f6843855ec1c14

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/430

We add losses in distillation by instantiating them in the distillation algorithm's init and then running them during the forward pass.

However this has some issues:
* the losses are not registered as a module in the model since they we organize them as a list of layerlossmetadata => this means that things like AMP do not behave as expected
* the losses are not on the same device as the rest of the model since they are created potentially after the model is moved to a new device

This diff solves both of these issues by including a helper function that registers and moves the losses to the same device as the model. `register_layer_losses_and_to_device` takes as input `List[LayerLossMetadata]`, moves the losses to the same device as the model and then registers these losses to the model.

Differential Revision: D41296932

fbshipit-source-id: ae7ae0847bce1b5cc481d838b9cae69cea424f25

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/429

Add a teacher type called `no_teacher` which can be specified by the user in the case they ignore the teacher (e.g., domain adaptation). Building the teacher just returns a noop (`nn.Identity`)

Differential Revision: D40971788

fbshipit-source-id: fc49ac44224c92806a7be253eefb8454305814eb

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/433

Distillation uses a module called `CachedLayer` to record the outputs of a layer to an internal dict. This dict is typically initialized by the object itself and any value is overwritten every time the model runs.

However, sometimes we need more than one output run of the layer (e.g., domain adaptation => we run the model on real, then synthetic data and need to use both outputs).

This diff adds a helper to set externally set the cache dict of a model. In other words, we can run `set_cache_dict` on some model to change the dict used by all `CachedLayer` in the model. This allows us to run the model and record some outputs, then change the cache dict and rerun the model to save different outputs.

Differential Revision: D40970577

fbshipit-source-id: 49cb851af49ae193d0c8ac9218e02fdaf4e6587b

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/421

Add some reasonable defaults when running knowledge distillation
* get_default_kd_image_classification_layer_losses => returns cross entropy loss on the output of the student classification layer and the teacher output (this is what the imagenet distillation uses)
* DefaultLossCombiner => simple function to multiply the losses by some weights

Unsure if these should go in `distillation.py` or a separate place (e.g., defaults or classification)

Reviewed By: chihyaoma

Differential Revision: D40330718

fbshipit-source-id: 5887566d88e3a96d01aca133c51041126b2692cc

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/420

Adds knowledge distillation as a generic algorithm that can be used by various projects.

If eval, the algorithm just returns the result of the student model.

If training, the algorithm feeds the input into both the student and teacher model. The user provides a list of `LayerLossMetadata` that provides the layers and losses run on these layers. The algorithm uses dynamic mixin to record the outputs of the relevant layers and compute the losses after both models are run.

We provide student and teacher preprocessing as a placeholder before we support a more generic dataloader which can provide different inputs to the student and teacher (e.g., as of now, if you want to provide the teacher with a larger input then the dataloader should return a large input and the student preprocessing can downsample the input).

We add the following functions as part of the user customizable distillation helper:
* get_teacher => return a teacher that can be used directly by the KD algorithm
* get_layer_losses => return a list of `LayerLossMetadata` that provides the layers and losses
* get_preprocess_student_input => manipulate the output of the dataloader before passing to the student
* get_preprocess_teacher_input => manipulate the output of the dataloader before passing to the teacher
* get_combine_losses => since we may want to weight the student and distillation losses, return a function that can manipulate the loss_dict

Reviewed By: chihyaoma

Differential Revision: D40326412

fbshipit-source-id: 2fb0e818a7d5b120d62fb7aba314ff96cc7e10c5

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/419

This diff adds a metadata class `LayerLossMetadata` to help keep track of the losses we want to compute over layers. The class contains the type of loss, loss name, and layer names.

This diff adds a helper function to iterate over a list of `LayerLossMetadata` and return a dict containing the results.

Reviewed By: chihyaoma

Differential Revision: D40286564

fbshipit-source-id: b269dc63cc90a437ca279379d759c3106016327c

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/418

This diff adds a function that can be used to add `CachedLayers` to a model. Function iterates over named modules and dynamically mixes in `CachedLayer` to target modules.

This diff adds a function to remove the cached layers.

Reviewed By: Minione

Differential Revision: D40285806

fbshipit-source-id: 3137d19927d8fb9ec924a77c9085aea29fe94d5e

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/417

This diff adds a layer `CachedLayer` which is meant to be used with dynamic mixin. This layer runs the original module and clones the output into a dictionary provided by the user.

The main use case is in distillation where we dynamically mixin these layers to the layers that the user wants to compute various losses.

See subsequent diffs to get integration with distillation.

Reviewed By: Minione

Differential Revision: D40285573

fbshipit-source-id: 2058deff8b96f63aebd1e9b9933a5352b5197111

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/416

Distillation assumes teacher model has an attribute "device". Sometimes this attribute is actually a property (e.g., generalizedrcnn) but there is zero guarantee that it exists. We add a helper function to move the model to the device and add this attribute if needed.

Reviewed By: chihyaoma

Differential Revision: D40283954

fbshipit-source-id: 42921653eac8a79499e22edac29aa6aeac016e8a

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/415

The user can build a teacher by providing a trained config. However this model may have been trained using gpu whereas the user wants to load the model on cpu, this diff supports this use case by allowing the user to specify `cfg.DISTILLATION.TEACHER.DEVICE` as override.

Reviewed By: sstsai-adl

Differential Revision: D40125236

fbshipit-source-id: f1fd797a155e12b31bb7fcbc5e4997ee8eb23539

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/371

In a previous iteration of this diff, we were specifying the teacher model in the same config as the student model, something like:
```
# config.py
MODEL:
  FBNET_V2:
  ...
DISTILLATION:
  TEACHER:
    MODEL:
      FBNET_V2:
      ...
      WEIGHTS: /path/to/teacher/weights
...
```

This leads to some oddities in the code, like we have to have a default config that adds all the required keys in the distillation teacher model.

In this diff, we just let the user supply a teacher config (and optionally runner_name and overwrite opts) and use the supplied runner to build the model:
```
# new_config.py
MODEL:
  FBNET_V2:
...
DISTILLATION:
  TEACHER:
    CONFIG_FNAME: /path/to/teacher/config
    RUNNER_NAME:
...
```

This should make it very easy to specify the teacher as the user could potentially just reuse the trained_config generated in d2go.

Reviewed By: newstzpz

Differential Revision: D37640041

fbshipit-source-id: 088a636c96f98279c9a04e32d1674f703451aec3

Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/318

Reviewed By: mcimpoi

Differential Revision: D37501246

fbshipit-source-id: 6dbe5dcbaf7454f451d4a3bb3fa2d856cc87d5cc

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/299

This implements the first iteration of generalized distillation in D2 (https://github.com/facebookresearch/d2go/commit/87374efb134e539090e0b5c476809dc35bf6aedb)Go. The functionality is separated into the following:

=> Adding distillation functionality without user changing their meta architecture:

class DistillationModelingHook
* This is an implementation detail that we hide from the user.
* We use dynamic mixin to specify additional functionality to the user's model. In this way, the original (student) model retains all attributes but the mixin class will override the forward (and provide more functionality like teacher updates).
* Teacher build currently only supports loading a torchscript model, pytorch compatiblity in later diffs

=> Implementing distillation methods

class DistillationAlgorithm
* The user can use some default algorithm (e.g., LabelDistillation) or create their own. This is where we specify the overrided forward func of the model and any other distillation requirements (updating the weights of the teacher model).
* The basic LabelDistillation allows a user to use a teacher model during training to relabel the gt

=> User customization

class DistillationHelper
* This is what we expect the user to customize. As an example the user probably needs to write their own pseudo_labeler to take batched_inputs and relabel this with the teacher

Both DistillationHelper and DistillationAlgorithm use a registry so that users can add their customization in their own code and use these customizations by specifying in the config

Reviewed By: newstzpz

Differential Revision: D36708227

fbshipit-source-id: bc427c5d42d0c7ff4d839bf10782efac24dea107