Repository: truera/trulens
Issue: Explain semantic segmentation models via PyTorch backend
Root cause: Unimplemented Feature Gap

Hi team, I'm delighted to explore this library for finding notorious issues in our trained models.

I'm trying to apply TruLens to various tasks (CV - segmentation, localisation; NLP - classification, generation etc.) and got stuck in the initial stages itself.

For started, I have tried to load pre-trained semantic/instance segmentation models and couldn't get them to work (after spending few cycles on debugging). 

A minute change in the [intro_demo_pytorch.ipynb](https://github.com/truera/trulens/blob/master/notebooks/intro_demo_pytorch.ipynb) at the line:

`pytorch_model = models.segmentation.fcn_resnet50(pretrained=True)`

throws the following error while calculating the _InputAttribution()_:

```
ValueError                                Traceback (most recent call last)
<ipython-input-11-9b10bbb33c8d> in <module>()
      1 infl = InputAttribution(model)
----> 2 attrs_input = infl.attributions(x_pp)

3 frames
/usr/local/lib/python3.7/dist-packages/trulens/nn/attribution.py in attributions(self, *model_args, **model_kwargs)
    269             intervention=D,
    270             doi_cut=doi_cut)
--> 271 
    272         # Take the mean across the samples in the DoI.
    273         if isinstance(qoi_grads, DATA_CONTAINER_TYPE):

/usr/local/lib/python3.7/dist-packages/trulens/nn/models/pytorch.py in qoi_bprop(self, qoi, model_args, model_kwargs, doi_cut, to_cut, attribution_cut, intervention)
    461         grads_list = []
    462         for z in zs:
--> 463             z_flat = ModelWrapper._flatten(z)
    464             qoi_out = qoi(y)
    465 

/usr/local/lib/python3.7/dist-packages/trulens/nn/quantities.py in __call__(self, y)
    109 
    110     def __call__(self, y: TensorLike) -> TensorLike:
--> 111         self._assert_cut_contains_only_one_tensor(y)
    112 
    113         if self.activation is not None:

/usr/local/lib/python3.7/dist-packages/trulens/nn/quantities.py in _assert_cut_contains_only_one_tensor(self, x)
     76                 '`{}` expected to receive an instance of `Tensor`, but '
     77                 'received an instance of {}'.format(
---> 78                     self.__class__.__name__, type(x)))
     79 
     80 

ValueError: `MaxClassQoI` expected to receive an instance of `Tensor`, but received an instance of <class 'collections.OrderedDict'>
```

Can you share some guidance on triaging this error? I'm using the same input image and transforms.

Hey Ashish, this error is because of an issue with Tensorflow 2.6 interfaces. #46 should fix this. In the meanwhile, feel free to try Tensorflow 2.4 or 2.5.

Hi Shayak, this error originates from the PyTorch backend. Should this have any correlation with TF version?

Ah you're right, missed that fact that this is the pytorch notebook. Will take a look and get back soon. Thanks for trying this out! 

Hi Ashish,

I think the issue is that the output to the model is an `OrderedDict` rather than a tensor. TruLens can still work with models with other types of outputs, but you will need a custom quantity of interest (`QoI`) to handle them.

If you can provide a bit more information about the output of your segmentation model, I can maybe give some advice on how to create your quantity of interest.

Klas, I'm using the segmentation model from Torchvision i.e. _torchvision.models.segmentation.fcn_resnet50_
The layer's of the model are attached below. I'm not specifying the from/to layer slices while initializing the _InputAttributions_ class, hence I suppose it'll automatically consider the first and last layers as slices.

Its a pretty standard segmentation model. Let me know if you'd like me to share more details.

```'backbone_conv1':	Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
'backbone_bn1':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_relu':	ReLU(inplace=True)
'backbone_maxpool':	MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
'backbone_layer1_0_conv1':	Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_0_bn1':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_0_conv2':	Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer1_0_bn2':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_0_conv3':	Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_0_bn3':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_0_relu':	ReLU(inplace=True)
'backbone_layer1_0_downsample_0':	Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_0_downsample_1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_1_conv1':	Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_1_bn1':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_1_conv2':	Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer1_1_bn2':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_1_conv3':	Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_1_bn3':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_1_relu':	ReLU(inplace=True)
'backbone_layer1_2_conv1':	Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_2_bn1':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_2_conv2':	Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer1_2_bn2':	BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_2_conv3':	Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer1_2_bn3':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer1_2_relu':	ReLU(inplace=True)
'backbone_layer2_0_conv1':	Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_0_bn1':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_0_conv2':	Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
'backbone_layer2_0_bn2':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_0_conv3':	Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_0_bn3':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_0_relu':	ReLU(inplace=True)
'backbone_layer2_0_downsample_0':	Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
'backbone_layer2_0_downsample_1':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_1_conv1':	Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_1_bn1':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_1_conv2':	Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer2_1_bn2':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_1_conv3':	Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_1_bn3':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_1_relu':	ReLU(inplace=True)
'backbone_layer2_2_conv1':	Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_2_bn1':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_2_conv2':	Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer2_2_bn2':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_2_conv3':	Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_2_bn3':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_2_relu':	ReLU(inplace=True)
'backbone_layer2_3_conv1':	Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_3_bn1':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_3_conv2':	Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer2_3_bn2':	BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_3_conv3':	Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer2_3_bn3':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer2_3_relu':	ReLU(inplace=True)
'backbone_layer3_0_conv1':	Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_0_bn1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_0_conv2':	Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'backbone_layer3_0_bn2':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_0_conv3':	Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_0_bn3':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_0_relu':	ReLU(inplace=True)
'backbone_layer3_0_downsample_0':	Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_0_downsample_1':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_1_conv1':	Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_1_bn1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_1_conv2':	Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2), bias=False)
'backbone_layer3_1_bn2':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_1_conv3':	Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_1_bn3':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_1_relu':	ReLU(inplace=True)
'backbone_layer3_2_conv1':	Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_2_bn1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_2_conv2':	Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2), bias=False)
'backbone_layer3_2_bn2':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_2_conv3':	Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_2_bn3':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_2_relu':	ReLU(inplace=True)
'backbone_layer3_3_conv1':	Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_3_bn1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_3_conv2':	Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2), bias=False)
'backbone_layer3_3_bn2':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_3_conv3':	Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_3_bn3':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_3_relu':	ReLU(inplace=True)
'backbone_layer3_4_conv1':	Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_4_bn1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_4_conv2':	Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2), bias=False)
'backbone_layer3_4_bn2':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_4_conv3':	Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_4_bn3':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_4_relu':	ReLU(inplace=True)
'backbone_layer3_5_conv1':	Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_5_bn1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_5_conv2':	Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2), bias=False)
'backbone_layer3_5_bn2':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_5_conv3':	Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer3_5_bn3':	BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer3_5_relu':	ReLU(inplace=True)
'backbone_layer4_0_conv1':	Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_0_bn1':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_0_conv2':	Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2), bias=False)
'backbone_layer4_0_bn2':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_0_conv3':	Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_0_bn3':	BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_0_relu':	ReLU(inplace=True)
'backbone_layer4_0_downsample_0':	Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_0_downsample_1':	BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_1_conv1':	Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_1_bn1':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_1_conv2':	Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(4, 4), dilation=(4, 4), bias=False)
'backbone_layer4_1_bn2':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_1_conv3':	Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_1_bn3':	BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_1_relu':	ReLU(inplace=True)
'backbone_layer4_2_conv1':	Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_2_bn1':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_2_conv2':	Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(4, 4), dilation=(4, 4), bias=False)
'backbone_layer4_2_bn2':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_2_conv3':	Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
'backbone_layer4_2_bn3':	BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'backbone_layer4_2_relu':	ReLU(inplace=True)
'classifier_0':	Conv2d(2048, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'classifier_1':	BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'classifier_2':	ReLU()
'classifier_3':	Dropout(p=0.1, inplace=False)
'classifier_4':	Conv2d(512, 21, kernel_size=(1, 1), stride=(1, 1))
'aux_classifier_0':	Conv2d(1024, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
'aux_classifier_1':	BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
'aux_classifier_2':	ReLU()
'aux_classifier_3':	Dropout(p=0.1, inplace=False)
'aux_classifier_4':	Conv2d(256, 21, kernel_size=(1, 1), stride=(1, 1))
```

Sorry, I meant the output of the model, not the architecture. I'm not particularly familiar with `torchvision`, but looking into it, it seems that the output of `torchvision` models is an `OrderedDict` and not a `torch.Tensor`, which by default has a single key, `'out'`. If this is the case for your model as well, you could try the following fix:
```python
class TorchvisionMaxClassQoI(MaxClassQoI):
    def __call__(self, y):
        super().__call__(y['out']) 
```
Then when you create your attribution measure you can use this QoI:
```python
infl = InputAttribution(model, qoi=TorchvisionMaxClassQoI())
attrs_input = infl.attributions(x_pp)
```
This assumes the output of your model is a dictionary containing a key, `'out'`, representing the model's actual output. If the output is different from that, you'll want to adjust how you extract the output tensor accordingly.