Issue: Setting to torch format not working with torchvision and MNIST

Hi

I am trying to use `torchvision.transforms` to handle the transformation of the image data in the `mnist` dataset. Assume I have a `transform` variable which contains the `torchvision.transforms` object.

A snippet of what I am trying to do:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                np.array(examples["image"][example_idx], dtype=np.uint8)
            ))
        else:
            images.append(torch.tensor(np.array(examples["image"][example_idx], dtype=np.uint8)))
        labels.append(torch.tensor(examples["label"][example_idx]))
    output = {"label":labels, "image":images}
    return output

raw_dataset = load_dataset('mnist')
train_dataset = raw_dataset.map(prepare_features, batched=True, batch_size=10000)
train_dataset.set_format("torch",columns=["image","label"])
```

After this, I check the type of the following:
```python
print(type(train_dataset["train"]["label"]))
print(type(train_dataset["train"]["image"][0]))
```
This leads to the following output:

```python
<class 'torch.Tensor'>
<class 'list'>
```
I use `torch.utils.DataLoader` for batches, the type of `batch["train"]["image"]` is also `<class 'list'>`.

I don't understand why only the `label` is converted to a torch tensor, why does the image not get converted? How can I fix this issue?

Thanks,
Gunjan

EDIT:
I just checked the shapes, and the types, `batch[image]` is a actually a list of list of tensors. Shape is (1,28,2,28), where `batch_size` is 2. I don't understand why this is happening. Ideally it should be a tensor of shape (2,1,28,28).

EDIT 2:
Inside `prepare_train_features`, the shape of `images[0]` is `torch.Size([1,28,28])`, the conversion is working. However, the output of the `map` is a list of list of list of list.

Comments:
Adding to the previous information, I think `torch.utils.data.DataLoader` is doing some conversion. 
What I tried:
```python
train_dataset = load_dataset('mnist')
```
I don't use any `map` or `set_format` or any `transform`. I use this directly, and try to load batches using the `DataLoader` with batch size 2, I get an output like this for the `image`:

```
[[tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor([0, 0]), tensor...
```
For `label`, it works fine:
```
tensor([7, 6])
```
Note that I didn't specify conversion to torch tensors anywhere.

Basically, there are two problems here:
1. `dataset.map` doesn't return tensor type objects, even though it uses the transforms, the grayscale conversion in transform was done, but the output was lists only.
2.  The `DataLoader` performs its own conversion, which may be not desired.

I understand that we can't change `DataLoader` because it is a torch functionality, however, is there a way we can handle image data to allow using it with torch `DataLoader` and `torchvision` properly?

I think if the `image` was a torch tensor (N,H,W,C), or  a list of torch tensors (H,W,C), before it is passed to `DataLoader`, then we might not face this issue. 

What's the feature types of your new dataset after `.map` ?

Can you try with adding `features=` in the `.map` call in order to set the "image" feature type to `Array2D` ?
The default feature type is lists of lists, we've not implemented shape verification to use ArrayXD instead of nested lists yet

Hi @lhoestq

Raw feature types are like this:
```
Image:
<class 'list'> 60000  #(type, len)
<class 'list'> 28
<class 'list'> 28
<class 'int'>
Label:
<class 'list'> 60000
<class 'int'>
```
Inside the `prepare_feature` method with batch size 100000 , after processing, they are like this:

Inside Prepare Train Features
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After map, the feature type are like this:
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'float'>
Label:
<class 'list'> 60000
<class 'int'>
```

After dataloader with batch size 2, the batch features are like this:
```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```
<hr>

When I was setting the format of `train_dataset` to 'torch' after mapping  - 
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```

Corresponding DataLoader batch:
```
From DataLoader batch features
Image:
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

I will check with features and get back.





Hi @lhoestq

# Using Array3D
I tried this:
```python
features = datasets.Features({
            "image": datasets.Array3D(shape=(1,28,28),dtype="float32"),
            "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
        })
train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)
```
and it didn't fix the issue.

During the `prepare_train_features:
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After the `map`:

```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'float'>
Label:
<class 'list'> 60000
<class 'int'>
```
From the DataLoader batch:
```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```
It is the same as before.

---

Using `datasets.Sequence(datasets.Array2D(shape=(28,28),dtype="float32"))` gave an error during `map`:

```python
ArrowNotImplementedError                  Traceback (most recent call last)
<ipython-input-95-d28e69289084> in <module>()
      3             "label": datasets.features.ClassLabel(names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]),
      4         })
----> 5 train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)

15 frames
/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    446                     num_proc=num_proc,
    447                 )
--> 448                 for k, dataset in self.items()
    449             }
    450         )

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1307                 fn_kwargs=fn_kwargs,
   1308                 new_fingerprint=new_fingerprint,
-> 1309                 update_data=update_data,
   1310             )
   1311         else:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    202         }
    203         # apply actual function
--> 204         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    205         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    206         # re-apply format to the output

/usr/local/lib/python3.7/dist-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    335             # Call actual function
    336 
--> 337             out = func(self, *args, **kwargs)
    338 
    339             # Update fingerprint of in-place transforms + update in-place history of transforms

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, update_data)
   1580                     if update_data:
   1581                         batch = cast_to_python_objects(batch)
-> 1582                         writer.write_batch(batch)
   1583             if update_data:
   1584                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    274             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    275             typed_sequence_examples[col] = typed_sequence
--> 276         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    277         self.write_table(pa_table, writer_batch_size)
    278 

/usr/local/lib/python3.7/dist-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.asarray()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._handle_arrow_array_protocol()

/usr/local/lib/python3.7/dist-packages/datasets/arrow_writer.py in __arrow_array__(self, type)
     95                 out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
     96             else:
---> 97                 out = pa.array(self.data, type=type)
     98             if trying_type and out[0].as_py() != self.data[0]:
     99                 raise TypeError(

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib.array()

/usr/local/lib/python3.7/dist-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

/usr/local/lib/python3.7/dist-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowNotImplementedError: extension
```

# Convert raw tensors to torch format
Strangely, converting to torch tensors works perfectly on `raw_dataset`:
```python
raw_dataset.set_format('torch',columns=['image','label'])
```
Types:
```
Image:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```

Using this for transforms:
```python
def prepare_features(examples):
    images = []
    labels = []
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(
                examples["image"][example_idx].numpy()
            ))
        else:
            images.append(examples["image"][example_idx].numpy())
        labels.append(examples["label"][example_idx])
    output = {"label":labels, "image":images}
    return output
```

Inside `prepare_train_features`:
```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After `map`:
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```
DataLoader batch:

```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

---

## Using `torch` format:
```
Image:
<class 'list'> 60000
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```
DataLoader batches:

```
Image:
<class 'list'> 1
<class 'list'> 28
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

---
## Using the features - `Array3D`:

```
Image:
<class 'list'> 10000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'list'> 10000
<class 'torch.Tensor'>
```

After `map`:
```
Image:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 60000
<class 'torch.Tensor'>
```

After DataLoader `batch`:
```
Image:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'> 1
<class 'torch.Tensor'> 28
<class 'torch.Tensor'> 28
<class 'torch.Tensor'>
Label:
<class 'torch.Tensor'> 2
<class 'torch.Tensor'>
```

The last one works perfectly.

![image](https://user-images.githubusercontent.com/29076344/110491452-4cf09c00-8117-11eb-8a47-73bf3fc0c3dc.png)

I wonder why this worked, and others didn't.








Concluding, the way it works right now is:

1. Converting raw dataset to `torch` format.
2. Use the transform and apply using `map`, ensure the returned values are tensors. 
3. When mapping, use `features` with `image` being `Array3D` type.

What the dataset returns depends on the feature type.
For a feature type that is Sequence(Sequence(Sequence(Value("uint8")))), a dataset formatted as "torch" return lists of lists of tensors. This is because the lists lengths may vary.
For a feature type that is Array3D on the other hand it returns one tensor. This is because the size of the tensor is fixed and defined bu the Array3D type.

Okay, that makes sense.
Raw images are list of Array2D, hence we get a single tensor when `set_format` is used. But, why should I need to convert the raw images to `torch` format when `map` does this internally?

Using `Array3D` did not work with `map` when raw images weren't `set_format`ted to torch type.

I understand that `map` needs to know what kind of output tensors are expected, and thus converting the raw dataset to `torch` format is necessary. Closing the issue since it is resolved.