Chrombpnet model conversion

[TFHan]

Dear Jacob,

first thanks for all your work on the DNA model front! I am aiming. to transfer pre-trained ChromBPNet models to PyTorch using from_chrombpnet() to fully leverage the tangermeme space. While the converted model performs well for count predictions, I’ve encountered an issue with SHAP scores when using a wrapped count head model. Specifically, the SHAP outputs appear to have a sudden cutoff/levelling of around 1200 base pairs, while random background around 400bp seems to have highest scores and extend down to bp-position 0. Inserting relevant motifs near or beyond the cutoff position seems to cause attribution scores to shrink or disappear.

I’d greatly appreciate any insights into potential causes of this issue, as using the converted models would greatly save time by avoiding the need for retraining.

All the best Timm

[jmschrei]

Hi Timm

Sorry that you're encountering issues. Can you post some example images to make sure I fully understand the issue?

And can you post your code for loading and wrapping the pre-trained model?

[TFHan]

Thanks for the quick response, here are the specs I used to load the model, along with the attribution plots comparing the same regions from the TensorFlow pipeline and the interpreted regions using the converted model:

bias_model_path = ".../bias.h5"
accessibility_model_path = ".../chrombpnet_nobias.h5"
model_name = "converted_chrombpnet"

chrombpnet_model = ChromBPNet.from_chrombpnet(bias_model_path, accessibility_model_path, model_name)

class WrappedChromBPNet(torch.nn.Module):
    def __init__(self, chrombpnet_model, output_index=0):
        super(WrappedChromBPNet, self).__init__()
        self.chrombpnet_model = chrombpnet_model
        self.output_index = output_index

    def forward(self, X):
        outputs = self.chrombpnet_model(X)
        return outputs[self.output_index].squeeze(1) 

wrapped_model = WrappedChromBPNet(chrombpnet_model, output_index=0)

[jmschrei]

It looks like you're sort of calculating the profile attributions using the PyTorch model, but not using the profile attribution wrapper that handles scaling and everything else correctly. Do you mean to do this, or are you trying to get the count attributions, which would be output_idx=1?

[TFHan]

thank you , i switched to using the built in count wrapper.

wrapped_count_model = CountWrapper(chrombpnet_model)

The resulting attributions are now approximately proportional to the TensorFlow version, while differing in range. Convergence deltas are too high, in the range of 0.1 to 0.7 running on CPU. Does this imply the requirement of additional_nonlinear_ops / other adjustments ?