pytorch中torch.unsqueeze()函数与np.expand_dims()

提示：扫一扫查出行
【扫一扫了解最新限行尾号】
复制提示

numpy. expand_dims ( a , axis )

Expand the shape of an array.

Insert a new axis that will appear at the axis position in the expanded array shape.

Parameters:	a : array_like Input array. axis : int Position in the expanded axes where the new axis is placed.
Returns:	res : ndarray Output array. The number of dimensions is one greater than that of the input array.

Examples

>>> x = np.array([1,2])

>>> x.shape

(2,)

The following is equivalent to x[np.newaxis,:] or x[np.newaxis] :

>>> y = np.expand_dims(x, axis=0)

>>> y

array([[1, 2]])

>>> y.shape

(1, 2)

>>> y = np.expand_dims(x, axis=1)  # Equivalent to x[:,np.newaxis]

>>> y

array([[1],

       [2]])

>>> y.shape

(2, 1)

Note that some examples may use None instead of np.newaxis . These are the same objects:

>>> np.newaxis is None

True

torch. unsqueeze ( input , dim , out=None ) → Tensor

Returns a new tensor with a dimension of size one inserted at the specified position.

The returned tensor shares the same underlying data with this tensor.

A dim value within the range [-input.dim() - 1, input.dim() + 1) can be used. Negative dim will correspond to unsqueeze() applied at dim = dim + input.dim() + 1 .

Parameters:	input ( Tensor ) – the input tensor dim ( int ) – the index at which to insert the singleton dimension out ( Tensor , optional ) – the output tensor

Example:

>>> x = torch.tensor([1, 2, 3, 4])

>>> torch.unsqueeze(x, 0)

tensor([[ 1,  2,  3,  4]])

>>> torch.unsqueeze(x, 1)

tensor([[ 1],

        [ 2],

        [ 3],

        [ 4]])