Enable uncertainty quantification on newer models

[z-q-y]

Uncertainty quantification based on Gal et al.'s work [https://arxiv.org/pdf/1703.04977.pdf] has been implemented on many models since issues #1119 and #1211 . However, for many of the newer models, e.g. deepchem.models.torch_models.DMPNNModel, this has not been done yet:

File "C:\Users\env\lib\site-packages\deepchem\models\torch_models\torch_model.py", line 809, in predict_uncertainty
    results = self._predict(generator, [], True, None)
  File "C:\Users\env\lib\site-packages\deepchem\models\torch_models\torch_model.py", line 575, in _predict
    raise ValueError('This model cannot compute uncertainties')
ValueError: This model cannot compute uncertainties

In the meantime, I would like guidance to come up with an ad hoc function to calculate epistemic uncertainty via Monte Carlo dropout. My main difficulty is that I'm having trouble applying dropout at inference. For PyTorch models, I suppose it could be done by simply switching on training mode, but I'm not sure how.

[z-q-y]

For PyTorch models, I suppose it could be done by simply switching on training mode, but I'm not sure how.

I followed through on this and applied a monkey patch, basically removing self.model.eval() from _predict() in TorchModel at inference. However, this is definitely not scientifically sound, as there are potentially other layers like BatchNorm that behave as if they are in training.

class TorchModelNew(TorchModel):
    def _predict(self, generator,
                 transformers, uncertainty,
                 other_output_types):
        results = None
        variances = None
        if uncertainty and (other_output_types is not None):
            raise ValueError(
                'This model cannot compute uncertainties and other output types simultaneously. Please invoke one at a time.'
            )
        if 1 or uncertainty:
            """
            if self._variance_outputs is None or len(
                    self._variance_outputs) == 0:
                raise ValueError('This model cannot compute uncertainties')
            if len(self._variance_outputs) != len(self._prediction_outputs):
                raise ValueError(
                    'The number of variances must exactly match the number of outputs'
                )
            """
            self._ensure_built()
            self.model.train()
        else:
            self._ensure_built()
            self.model.eval()
        if other_output_types:
            if self._other_outputs is None or len(self._other_outputs) == 0:
                raise ValueError(
                    'This model cannot compute other outputs since no other output_types were specified.'
                )
        for batch in generator:
            inputs, labels, weights = batch
            inputs, _, _ = self._prepare_batch((inputs, None, None))

            # Invoke the model.
            if isinstance(inputs, list) and len(inputs) == 1:
                inputs = inputs[0]
            output_values = self.model(inputs)
            if isinstance(output_values, torch.Tensor):
                output_values = [output_values]
            output_values = [t.detach().cpu().numpy() for t in output_values]

            # Apply tranformers and record results.
            if uncertainty:
                var = [output_values[i] for i in self._variance_outputs]
                if variances is None:
                    variances = [var]
                else:
                    for i, t in enumerate(var):
                        variances[i].append(t)
            access_values = []
            if other_output_types:
                access_values += self._other_outputs
            elif self._prediction_outputs is not None:
                access_values += self._prediction_outputs

            if len(access_values) > 0:
                output_values = [output_values[i] for i in access_values]

            if len(transformers) > 0:
                if len(output_values) > 1:
                    raise ValueError(
                        "predict() does not support Transformers for models with multiple outputs."
                    )
                elif len(output_values) == 1:
                    output_values = [
                        undo_transforms(output_values[0], transformers)
                    ]
            if results is None:
                results = [[] for i in range(len(output_values))]
            for i, t in enumerate(output_values):
                results[i].append(t)

        # Concatenate arrays to create the final results.
        final_results = []
        final_variances = []
        if results is not None:
            for r in results:
                final_results.append(np.concatenate(r, axis=0))
        if uncertainty and variances is not None:
            for v in variances:
                final_variances.append(np.concatenate(v, axis=0))
            return zip(final_results, final_variances)
        if len(final_results) == 1:
            return final_results[0]
        else:
            return final_results

[rbharath]

@z-q-y This is a very worthwhile project but definitely a larger one. Can you come by OH sometime? Or if not, could you put together a full design doc for review? I can read through a proposed approach and give you feedback