python-3.x - 从 Numpy 数组中获取变量

Question

对于 Pytorch 的图像分类问题，我在数据中读取如下：

import scipy .io
emnist = scipy.io.loadmat(DATA_DIR + '/emnist-letters.mat')
data = emnist ['dataset']
X_train = data ['train'][0, 0]['images'][0, 0]
X_train = X_train.reshape((-1,28,28), order='F')

y_train = data ['train'][0, 0]['labels'][0, 0]

X_test = data ['test'][0, 0]['images'][0, 0]
X_test = X_test.reshape((-1,28,28), order = 'F')

y_test = data ['test'][0, 0]['labels'][0, 0]

我的目标是创建一个数据集，使用：

train_dataset = torch.utils.data.TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))

目前，当我运行 的实例时y_train，输出是一个数组：

y_train[0]
>>> array([23], dtype=uint8)

但是，我只想train_dataset包含数组内第 0 个索引处的数字（在本例中为 23），而不是整个数组。

如何更改我的代码，以便创建的 TensorDataset 仅包含数组的第一个元素y_train，而不是整个数组？

标签： python-3.xnumpyindexingpytorchtorch

Answer 1

您可以使用np.squeeze()删除大小为 1 的数据的维度。如果您要删除特定维度，请将其传递给squeeze.

import numpy as np

arr = np.random.randn(1, 2, 1, 3, 1)
arr.squeeze().shape # (2, 3)

arr.squeeze(2).shape # (1, 2, 3, 1)