Swin-Transformer-TF not working with generator

Question

Swin-Transformer-TF from https://github.com/rishigami/Swin-Transformer-TF/tree/main, failed when I use generator for train. I tried many generators, and many output of generators options, lists, Tensor's etc. everything is crashing.

import tensorflow as tf
import numpy as np
!git clone https://github.com/rishigami/Swin-Transformer-TF.git
import sys
sys.path.append('./Swin-Transformer-TF')
from swintransformer import SwinTransformer

create the model:

img_adjust_layer = tf.keras.layers.Lambda(lambda data: tf.keras.applications.imagenet_utils.preprocess_input(tf.cast(data, 
                                                                                                                     tf.float32), 
                                                                                                             mode="torch"), 
                                          input_shape=[224, 224,3])
pretrained_model = SwinTransformer('swin_tiny_224', num_classes=2, 
                                   include_top=True, pretrained=0, 
                                   use_tpu=0)

model = tf.keras.Sequential([
    img_adjust_layer,
    pretrained_model,
    tf.keras.layers.Dense(2, activation='softmax')
])
    
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5, epsilon=1e-8),
    loss = 'categorical_crossentropy',
    metrics=['accuracy']
)
model.summary()

Train the model, without generetor, works.

x_train = np.zeros((20,224,224,3))
y_train = np.zeros((20,2))
model.fit(x_train, y_train,
         steps_per_epoch=2, epochs=1)

Adding a generator (and I tried a few), crashes. The shape is read as (None,None,None,3):

def __data_generation():
    for i in range(3000):
      yield np.zeros((20,224,224,3)),np.zeros((20,2))
gen=__data_generation()
x,y=next(gen)
print(x.shape,y.shape)
model.fit(gen, epochs=1, steps_per_epoch=4)



AssertionError                            Traceback (most recent call last)
<ipython-input-7-fb53f3a3d1ae> in <cell line: 7>()
      5 x,y=next(gen)
      6 print(x.shape,y.shape)
----> 7 model.fit(gen, epochs=1, steps_per_epoch=4)

4 frames
/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

/usr/local/lib/python3.10/dist-packages/keras/engine/training.py in tf__train_function(iterator)
     13                 try:
     14                     do_return = True
---> 15                     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16                 except:
     17                     do_return = False

/content/./Swin-Transformer-TF/swintransformer/model.py in tf__call(self, x)
      8                 do_return = False
      9                 retval_ = ag__.UndefinedReturnValue()
---> 10                 x = ag__.converted_call(ag__.ld(self).forward_features, (ag__.ld(x),), None, fscope)
     11 
     12                 def get_state():

/content/./Swin-Transformer-TF/swintransformer/model.py in tf__forward_features(self, x)
      8                 do_return = False
      9                 retval_ = ag__.UndefinedReturnValue()
---> 10                 x = ag__.converted_call(ag__.ld(self).patch_embed, (ag__.ld(x),), None, fscope)
     11 
     12                 def get_state():

/content/./Swin-Transformer-TF/swintransformer/model.py in tf__call(self, x)
      9                 retval_ = ag__.UndefinedReturnValue()
     10                 (B, H, W, C) = ag__.converted_call(ag__.converted_call(ag__.ld(x).get_shape, (), None, fscope).as_list, (), None, fscope)
---> 11                 assert ag__.and_(lambda : ag__.ld(H) == ag__.ld(self).img_size[0], lambda : ag__.ld(W) == ag__.ld(self).img_size[1]), f"Input image size ({ag__.ld(H)}*{ag__.ld(W)}) doesn't match model ({ag__.ld(self).img_size[0]}*{ag__.ld(self).img_size[1]})."
     12                 x = ag__.converted_call(ag__.ld(self).proj, (ag__.ld(x),), None, fscope)
     13                 x = ag__.converted_call(ag__.ld(tf).reshape, (ag__.ld(x),), dict(shape=[-1, ag__.ld(H) // ag__.ld(self).patch_size[0] * (ag__.ld(W) // ag__.ld(self).patch_size[0]), ag__.ld(self).embed_dim]), fscope)

AssertionError: in user code:

    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1284, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1268, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1249, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1050, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_filegq78ow1y.py", line 10, in tf__call
        x = ag__.converted_call(ag__.ld(self).forward_features, (ag__.ld(x),), None, fscope)
    File "/tmp/__autograph_generated_filehkbim9xc.py", line 10, in tf__forward_features
        x = ag__.converted_call(ag__.ld(self).patch_embed, (ag__.ld(x),), None, fscope)
    File "/tmp/__autograph_generated_file3nqntnhc.py", line 11, in tf__call
        assert ag__.and_(lambda : ag__.ld(H) == ag__.ld(self).img_size[0], lambda : ag__.ld(W) == ag__.ld(self).img_size[1]), f"Input image size ({ag__.ld(H)}*{ag__.ld(W)}) doesn't match model ({ag__.ld(self).img_size[0]}*{ag__.ld(self).img_size[1]})."

    AssertionError: Exception encountered when calling layer 'swin_tiny_224' (type SwinTransformerModel).
    
    in user code:
    
        File "/content/./Swin-Transformer-TF/swintransformer/model.py", line 422, in call  *
            x = self.forward_features(x)
        File "/content/./Swin-Transformer-TF/swintransformer/model.py", line 411, in forward_features  *
            x = self.patch_embed(x)
        File "/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
        File "/tmp/__autograph_generated_file3nqntnhc.py", line 11, in tf__call
            assert ag__.and_(lambda : ag__.ld(H) == ag__.ld(self).img_size[0], lambda : ag__.ld(W) == ag__.ld(self).img_size[1]), f"Input image size ({ag__.ld(H)}*{ag__.ld(W)}) doesn't match model ({ag__.ld(self).img_size[0]}*{ag__.ld(self).img_size[1]})."
    
        AssertionError: Exception encountered when calling layer 'patch_embed' (type PatchEmbed).
        
        in user code:
        
            File "/content/./Swin-Transformer-TF/swintransformer/model.py", line 336, in call  *
                assert H == self.img_size[0] and W == self.img_size[1],             f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        
            AssertionError: Input image size (None*None) doesn't match model (224*224).
        
        
        Call arguments received by layer 'patch_embed' (type PatchEmbed):
          • x=tf.Tensor(shape=(None, None, None, 3), dtype=float32)
    
    
    Call arguments received by layer 'swin_tiny_224' (type SwinTransformerModel):
      • x=tf.Tensor(shape=(None, None, None, 3), dtype=float32)

Answer 1

tf.data.Dataset.from_generator should help in this case, when working with generators. Modified code snippet:

import tensorflow as tf 
import numpy as np 


def __data_generation():
    for i in range(20):
      yield np.zeros((20,224,224,3)),np.zeros((20,2))

# Directly using the generator doesn't work
# gen=__data_generation()
# x,y=next(gen)
# print(x.shape,y.shape)

# Try wrapping it in tf.data pipeline
ds = tf.data.Dataset.from_generator(
    __data_generation, 
    output_signature=(
        tf.TensorSpec(shape=(None,224,224,3), dtype=tf.float32), 
        tf.TensorSpec(shape=(None,2,), dtype=tf.float32),
    )
)
x, y = next(iter(ds))
print("Shapes: ", x.shape, y.shape)  # Shapes:  (20, 224, 224, 3) (20, 2)

model.fit(ds, epochs=1, steps_per_epoch=4)

Moreover, if the generator outputs single image-label pairs (shape (224,224,3), (2,)) instead of batching them already (e.g. shape (20,224,224,3), (20,2) where batch size is 20), we can use ds.batch() to set the batch size before calling fit():

def __data_generation():
    for i in range(20):
      yield np.zeros((224,224,3)),np.zeros((2,))

# Try wrapping it in tf.data pipeline
ds = tf.data.Dataset.from_generator(
    __data_generation, 
    output_signature=(
        tf.TensorSpec(shape=(224,224,3), dtype=tf.float32), 
        tf.TensorSpec(shape=(2,), dtype=tf.float32),
    )
)
ds = ds.batch(2)  # batch size = 2

model.fit(ds, epochs=1, steps_per_epoch=4)

More guidelines to tf.data can be found in the documentation.