adapter that handles torch utils data dataloader def initself dataloader import torch if not isinstancedataloader torch utils data dataloader raise valueerror fexpected argument dataloader to be an instance of ftorch utils data dataloader received dataloader self dataloader dataloader self batchsize dataloader batchsize self size lendataloader self partialbatchsize lendataloader dataset self batchsize def getnumpyiteratorself for batch in self dataloader yield tupletree mapstructurelambda x x cpu numpy batch def gettorchdataloaderself return self dataloader def gettfdatasetself from keras utils moduleutils import tensorflow as tf outputsignature self peekandgettensorspec return tf data dataset fromgenerator self getnumpyiterator outputsignatureoutputsignature def peekandgettensorspecself from keras utils moduleutils import tensorflow as tf batchdata nextiterself dataloader def gettensorspecx shape x shape if lenshape 1 raise valueerror when passing a pytorch dataloader to a keras model the arrays returned by the generator must be at least rank 1 received fx of rank lenx shape shape listshape shape0 none the batch size is not guaranteed to be static no easy way to get string representation of dtype in torch todo figure out a better way to achieve this dtype strx dtype replacetorch return tf tensorspecshapeshape dtypedtype return tupletree mapstructuregettensorspec batchdata property def numbatchesself return self size property def batchsizeself return self batchsize property def haspartialbatchself if self partialbatchsize return self partialbatchsize 0 else return none property def partialbatchsizeself return self partialbatchsize adapter that handles torch utils data dataloader the batch size is not guaranteed to be static no easy way to get string representation of dtype in torch todo figure out a better way to achieve this
import tree

from keras.trainers.data_adapters.data_adapter import DataAdapter


class TorchDataLoaderAdapter(DataAdapter):
    

    def __init__(self, dataloader):
        import torch

        if not isinstance(dataloader, torch.utils.data.DataLoader):
            raise ValueError(
                f"Expected argument `dataloader` to be an instance of"
                f"`torch.utils.data.DataLoader`. Received: {dataloader}"
            )

        self._dataloader = dataloader
        self._batch_size = dataloader.batch_size
        self._size = len(dataloader)
        self._partial_batch_size = len(dataloader.dataset) % self._batch_size

    def get_numpy_iterator(self):
        for batch in self._dataloader:
            yield tuple(tree.map_structure(lambda x: x.cpu().numpy(), batch))

    def get_torch_dataloader(self):
        return self._dataloader

    def get_tf_dataset(self):
        from keras.utils.module_utils import tensorflow as tf

        output_signature = self.peek_and_get_tensor_spec()
        return tf.data.Dataset.from_generator(
            self.get_numpy_iterator,
            output_signature=output_signature,
        )

    def peek_and_get_tensor_spec(self):
        from keras.utils.module_utils import tensorflow as tf

        batch_data = next(iter(self._dataloader))

        def get_tensor_spec(x):
            shape = x.shape
            if len(shape) < 1:
                raise ValueError(
                    "When passing a Pytorch DataLoader to a Keras model, "
                    "the arrays returned by the generator "
                    "must be at least rank 1. Received: "
                    f"{x} of rank {len(x.shape)}"
                )
            shape = list(shape)
            shape[0] = None  

            
            
            dtype = str(x.dtype).replace("torch.", "")
            return tf.TensorSpec(shape=shape, dtype=dtype)

        return tuple(tree.map_structure(get_tensor_spec, batch_data))

    @property
    def num_batches(self):
        return self._size

    @property
    def batch_size(self):
        return self._batch_size

    @property
    def has_partial_batch(self):
        if self._partial_batch_size:
            return self._partial_batch_size > 0
        else:
            return None

    @property
    def partial_batch_size(self):
        return self._partial_batch_size