# Block masking in TensorFlow

**URL:** <https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727>\
**Category:** General Discussion\
**Tags:** help\_request\
**Created:** [March 14, 2022, 10:40am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727 "2022-03-14T10:40:49Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 14, 2022, 10:40am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/1 "2022-03-14T10:40:49Z")

</div>

I’m trying to implement the masking generation function for BEiT:

> <https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py>

The part I am struggling with is the assignment of EagerTensors.

I have consulted references that show how to approach such assignments, but this one does not seem to fit them.

Any particular approaches I should try out or look into for this case?

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 14, 2022, 11:26am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/3 "2022-03-14T11:26:11Z")

</div>

Is every single masking patch random inside the single image there?

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 14, 2022, 11:28am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/4 "2022-03-14T11:28:02Z")

</div>

P.s. I was looking at:

> **[keras\_cv\_attention\_models/keras\_cv\_attention\_models/beit at main ·...](https://github.com/leondgarse/keras_cv_attention_models/tree/main/keras_cv_attention_models/beit)**
>
> Keras beit,caformer,CMT,CoAtNet,convnext,davit,dino,efficientdet,edgenext,efficientformer,efficientnet,eva,fasternet,fastervit,fastvit,flexivit,gcvit,ghostnet,gpvit,hornet,hiera,iformer,inceptionne...

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 14, 2022, 11:45am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/5 "2022-03-14T11:45:56Z")

</div>

A single mask can be applied to a batch too.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 14, 2022, 11:46am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/6 "2022-03-14T11:46:13Z")

</div>

Thanks for sharing this. Will take a look.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 15, 2022, 8:29am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/8 "2022-03-15T08:29:43Z")

</div>

@Bhack there’s actually no masking involved in the link you sent.

So, the question is pretty much still open.

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 15, 2022, 10:20am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/9 "2022-03-15T10:20:24Z")

</div>

Yes as they are just reloading Microsoft weights. So no train protocol there.

What Is your specific issue? Isn’t just the standard  
image tokenization in many visual transformer where some token are masked?

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 15, 2022, 10:44am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/10 "2022-03-15T10:44:11Z")

</div>

> What Is your specific issue? Isn’t just the standard  
> image tokenization in many visual transformer where some token are masked?

My issue is in the block-wise masking strategy where apparently tensor assignment is needed (refer to my initial post). Had it been randomized, it would have been easier and we implemented that a while back ([here](https://keras.io/examples/vision/masked_image_modeling/)).

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 15, 2022, 12:21pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/11 "2022-03-15T12:21:33Z")

</div>

To exactly mimic that impl are you looking for slice assigment?

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 15, 2022, 12:41pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/12 "2022-03-15T12:41:10Z")

</div>

Yes. Please take note of this part before sharing existing references:

> I have consulted references that show how to approach such assignments, but this one does not seem to fit them.

If there’s no way other than doing something like [this](https://towardsdatascience.com/how-to-replace-values-by-index-in-a-tensor-with-tensorflow-2-0-510994fe6c5f), then it’s a different choice.

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 15, 2022, 1:05pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/13 "2022-03-15T13:05:15Z")

</div>

Oh, in that case historically we are full of slice assignment tickets. Just to mention a few still open:

> <https://github.com/tensorflow/tensorflow/issues/40605>
>
> \*\*System information\*\*
> \- TensorFlow version (you are using): 2.2
> \- Are you wil…ling to contribute it (Yes/No): Yes
> 
> \*\*Describe the feature and the current behavior/state.\*\*
> 
> I would like to have slice assignment for Tensor objects in TensorFlow.
> The code I would like to write is:
> 
> \`\`\`python
> import tensorflow as tf
> 
> a = tf.constant(\[1, 2, 4, 5, 7, 3, 2, 6,\])
> indices = tf.constant(\[3, 4\], dtype=tf.int32)
> 
> a\[indices\] += 1
> \`\`\`
> 
> Of course it's a simplistic example and doesn't cover everything I want to do (I would use it in more complex functions not with constants), and I am happy to make it more complex if necessary.
> 
> Currently this code gives the error:
> \`\`\`
> TypeError: Only integers, slices (\`:\`), ellipsis (\`...\`), tf.newaxis (\`None\`) and scalar tf.int32/tf.int64 tensors are valid indices, got \<tf.Tensor: shape=(2,), dtype=int32, numpy=array(\[3, 4\], dtype=int32)\>
> \`\`\`
> 
> \*\*Will this change the current api? How?\*\*
> 
> I guess this is a change of API since it introduces a new functionality.
> 
> \*\*Who will benefit with this feature?\*\*
> 
> A lot of people have been asking for this feature for example in this GitHub issues:
> \- https://github.com/tensorflow/tensorflow/issues/14132#issuecomment-483002522
> \- https://github.com/tensorflow/tensorflow/issues/33131
> 
> These issues have unfortunately been closed because some workarounds for specific use-cases have been found (ones where the slicing is fixed and you can use \[masking\](https://github.com/tensorflow/tensorflow/issues/14132#issuecomment-483002522) or \[TensorArrays\](https://github.com/tensorflow/tensorflow/issues/14132#issuecomment-487643287)).
> Some other issues deal with \`Variable\`s which is not what I am talking about here. \[Some workarounds do exist\](https://stackoverflow.com/a/62202181/4332585) involving \`Variable\` but they seem hacky.
> 
> I will personally benefit from it, in the multiple places where I now use \`tensor\_scatter\_nd\_add\` or \`tensor\_scatter\_nd\_update\`, which is solution that always works but is very difficult to write and very slow:
> 
> \- \[for a wavelet-based neural network, called MWCNN\](https://github.com/zaccharieramzi/tf-mwcnn/blob/master/mwcnn.py#L106-L110);
> \- \[for non-uniform fast fourier transform\](https://github.com/zaccharieramzi/tfkbnufft/blob/master/tfkbnufft/nufft/interp\_functions.py#L151);
> \- \[for sensitivity map extraction when doing MRI reconstruction with TensorFlow neural networks\](https://github.com/zaccharieramzi/fastmri-reproducible-benchmark/blob/master/fastmri\_recon/data/utils/multicoil/smap\_extract.py#L27-L35).
> 
> \*\*Any Other info.\*\*
> 
> The \`tensor\_scatter\_nd\_\*\` alternative might seem like a viable solution, but it suffers from 2 drawbacks that I consider huge:
> \- It is very difficult to write. It is actually so difficult, I decided to make a package that would alleviate this difficulty by having the different slicing possibilities unit tested: \[tf-slice-assign\](https://github.com/zaccharieramzi/tf-slice-assign).
> \- It is very slow. I made a \[benchmark notebook\](https://colab.research.google.com/drive/1gEjha7h1mhQkFwULS9MAU0bWQfzfEALY?usp=sharing) vs \`pytorch\` for slice assignment add. You can see that on GPU, using \`tensor\_scatter\_nd\_add\` is 10 times slower than slice assignment in \`pytorch\` and 20 times slower on CPU. For a practical example, it means that my \`tfkbnufft\` (for non-uniform fast fourier transform) package is 30 times slower than its \[torch counterpart\](https://github.com/mmuckley/torchkbnufft#computation-speed) which I translated. This currently removes the possibility of training neural networks using the non-uniform fourier transform in TensorFlow.

> <https://github.com/tensorflow/tensorflow/issues/33131>
>
> as in numpy or pytorch ,we can do someting like this, but how to do it with tf2.…0.
> the following code will raise exception as :
> 
> \`'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment\`
> prediction\[:,:,0\]=tf.math.sigmoid(prediction\[:,:,0\])

> <https://github.com/tensorflow/tensorflow/issues/26839>
>
> I was wondering if it is possible to implement Numpy like slicing annd updating …a\[1:10,2:20....\] in Tensorflow. It would make life much easier. Right now it code just gets bigger and uglier and bug prone.

I’ve not checked the paper in details on what kind of index is going to be selected to execute the masking. Cannot be covered by `tf.tensor_scatter_nd_update` after populating these indexes?

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 15, 2022, 1:18pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/14 "2022-03-15T13:18:30Z")

</div>

The indexing conditions are in the source code I provided.

If you know a way around with scatter, do you mind providing a minimal working code.

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 15, 2022, 1:27pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/15 "2022-03-15T13:27:07Z")

</div>

E.g. I think that embedding in the Hugginface transformers library, also if it is using Pytorch ops, is not going to require/use the slice assignment:

> <https://github.com/huggingface/transformers/blob/main/src/transformers/models/beit/modeling_beit.py#L138-L214>

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 15, 2022, 4:51pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/16 "2022-03-15T16:51:14Z")

</div>

I think you’re mistaken then.

`bool_masked_pos` in the `forward()` is nothing but the output the mask yielded by the class I showed in my initial post.

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 15, 2022, 9:09pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/17 "2022-03-15T21:09:03Z")

</div>

> [@Bhack](#):
>
> `tensor_scatter_nd_update`

It is true `bool_masked_pos` is only the “application” of the masking but then ownership to prepare the mask it is still to the external the caller.

I don’t see all the details that are in reference implementation in the paper but with the concrete reference implementation you shared, with all these `attemps`, `conditional loops` etc, you could try to use a `tf.variable` to mimic that implementation but probably you will need to refactor it more in graph mode/`tf.function`:

> **[Google Colab](https://colab.research.google.com/gist/bhack/2d2cf17707d904157e9d651c448d8174/untitled269.ipynb)**

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 16, 2022, 1:04am UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/18 "2022-03-16T01:04:28Z")

</div>

> [@Bhack](#):
>
> It is true `bool_masked_pos` is only the “application” of the masking but then ownership to prepare the mask it is still to the external the caller.

Absolutely. And in case no reference implementations are available I guess the implementation done by the actual author comes to the rescue.

> [@Bhack](#):
>
> I don’t see all the details that are in reference implementation in the paper

There isn’t much about it in the paper apart from the figure on block-wise masking which is why the original implementation is an important reference point.

Thanks for sharing your implementation. Will check it out.

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 17, 2022, 1:58pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/19 "2022-03-17T13:58:59Z")

</div>

> [@Sayak\_Paul](#):
>
> Thanks for sharing your implementation. Will check it out.

Having a `tf.function/graph` version It is quite trivial with few changes/substitution with TF ops.

But a `jit_compile=True` version it will require a new design and probably some compromises.

Let me know if you have a `jit_comile=True` version.

---

<div class="post-metadata">

**Author:** ![Sayak\_Paul](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/sayak_paul/32/13002_2.png) [@Sayak\_Paul](https://discuss.ai.google.dev/u/Sayak_Paul)\
**Post date:** [March 17, 2022, 2:19pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/20 "2022-03-17T14:19:27Z")

</div>

What’s trivial for you may not be trivial to someone else 🙂

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 17, 2022, 2:25pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/21 "2022-03-17T14:25:45Z")

</div>

Let me know when you have the same I’ve posted but with TF instead of numpy ops.  
I will help you to make the required changes for `tf.function`.

---

<div class="post-metadata">

**Author:** ![Bhack](https://d1dlmcr85iqnpo.cloudfront.net/user_avatar/discuss.ai.google.dev/bhack/32/6388_2.png) [@Bhack](https://discuss.ai.google.dev/u/Bhack)\
**Post date:** [March 17, 2022, 10:23pm UTC](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727/22 "2022-03-17T22:23:22Z")

</div>

This is already working with `tf.function` with minimal changes

```python
    @tf.function
    def _mask(self, mask, max_mask_patches):
        delta = 0
        for attempt in tf.range(10):
            target_area = random.uniform(self.min_num_patches, max_mask_patches)
            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))
            if w < self.width and h < self.height:
                top = random.randint(0, int(self.height - h))
                left = random.randint(0, int(self.width - w))
                num_masked = tf.math.count_nonzero(mask[top: top + h, left: left + w])
                # Overlap
                if 0 < h * w - num_masked and h * w - num_masked <= max_mask_patches:
                    for i in range(top, top + h):
                        for j in range(left, left + w):
                            if mask[i, j] == 0:
                                mask[i, j].assign(1)
                                delta += 1
                if delta > 0:
                    break
        return delta

```

[Next page](https://discuss.ai.google.dev/t/block-masking-in-tensorflow/28727.md?page=2)