numpyとtorchのdimension拡大と縮小

背景

numpyとpytorchには次元拡張と削減する方法は機械学習には頻繁に使われてます。今回は軽くそれを説明いたします。
次元拡張
np.expand_dims
torch.unsqueeze

次元縮小
np.squeeze
torch.squeeze

1.numpy

1-1.np.expand_dims

次元拡張

テストコード

a = np.array([
    [1, 3, 5],
    [1, "A", 3],
    [5, 1, "F"],
    [4, "G", "S"],
    ["S", "G", 2]
])


print(a)
print("np.expand_dims(a, 0)")
b = np.expand_dims(a, 0)
print(b)


print("np.expand_dims(a, 1)")
b = np.expand_dims(a, 1)
print(b)


print("np.expand_dims(a, 2)")
b = np.expand_dims(a, 2)
print(b)

output

[['1' '3' '5']
 ['1' 'A' '3']
 ['5' '1' 'F']
 ['4' 'G' 'S']
 ['S' 'G' '2']]

np.expand_dims(a, 0)
[[['1' '3' '5']
  ['1' 'A' '3']
  ['5' '1' 'F']
  ['4' 'G' 'S']
  ['S' 'G' '2']]]

np.expand_dims(a, 1)
[[['1' '3' '5']]

 [['1' 'A' '3']]

 [['5' '1' 'F']]

 [['4' 'G' 'S']]

 [['S' 'G' '2']]]

np.expand_dims(a, 2)
[[['1']
  ['3']
  ['5']]

 [['1']
  ['A']
  ['3']]

 [['5']
  ['1']
  ['F']]

 [['4']
  ['G']
  ['S']]

 [['S']
  ['G']
  ['2']]]

1-2.np.squeeze

次元縮小

テストコード

a = np.array([
    [1, 3, 5],
    [1, "A", 3],
    [5, 1, "F"],
    [4, "G", "S"],
    ["S", "G", 2]
])

print("a")
print(a)
print()
print("np.expand_dims(a,2)")
b = np.expand_dims(a,2)
print(b)

print()
print("c = np.squeeze(b,2)")
c = np.squeeze(b,2)
print(c)

print()
print("d = np.expand_dims(b,0)")
d = np.expand_dims(b,0)
print(d)

print()
print("e = np.squeeze(d,3)")
e = np.squeeze(d,3)
print(e)

output

a
[['1' '3' '5']
 ['1' 'A' '3']
 ['5' '1' 'F']
 ['4' 'G' 'S']
 ['S' 'G' '2']]

np.expand_dims(a,2)
[[['1']
  ['3']
  ['5']]

 [['1']
  ['A']
  ['3']]

 [['5']
  ['1']
  ['F']]

 [['4']
  ['G']
  ['S']]

 [['S']
  ['G']
  ['2']]]

c = np.squeeze(b,2)
[['1' '3' '5']
 ['1' 'A' '3']
 ['5' '1' 'F']
 ['4' 'G' 'S']
 ['S' 'G' '2']]

d = np.expand_dims(b,0)
[[[['1']
   ['3']
   ['5']]

  [['1']
   ['A']
   ['3']]

  [['5']
   ['1']
   ['F']]

  [['4']
   ['G']
   ['S']]

  [['S']
   ['G']
   ['2']]]]

e = np.squeeze(d,3)
[[['1' '3' '5']
  ['1' 'A' '3']
  ['5' '1' 'F']
  ['4' 'G' 'S']
  ['S' 'G' '2']]]

2.torch

2-1.torch.unsqueeze

次元拡大

テストコード

import torch
a = torch.tensor([
    [1, 3, 5],
    [1, 2, 3],
    [5, 1, 3],
    [4, 2, 7],
    [98, 11, 2]
])

c0 = torch.unsqueeze(a, 0)
print(c0)

c1 = torch.unsqueeze(a, 1)
print(c1)

c2 = torch.unsqueeze(a, 2)
print(c2)

output

tensor([[[ 1,  3,  5],
         [ 1,  2,  3],
         [ 5,  1,  3],
         [ 4,  2,  7],
         [98, 11,  2]]])
tensor([[[ 1,  3,  5]],

        [[ 1,  2,  3]],

        [[ 5,  1,  3]],

        [[ 4,  2,  7]],

        [[98, 11,  2]]])
tensor([[[ 1],
         [ 3],
         [ 5]],

        [[ 1],
         [ 2],
         [ 3]],

        [[ 5],
         [ 1],
         [ 3]],

        [[ 4],
         [ 2],
         [ 7]],

        [[98],
         [11],
         [ 2]]])

2-2.torch.squeeze

次元縮小

テストコード

import torch
a = torch.tensor([
    [[1], [3], [5]],
    [[1], [2], [3]],
    [[5], [1], [3]],
    [[4], [2], [7]],
    [[98], [11], [2]]
])

print("torch.squeeze(a, 0)")
c0 = torch.squeeze(a, 0)
print(c0)

print()
print("torch.squeeze(a, 1)")
c1 = torch.squeeze(a, 1)
print(c1)

print()
print("torch.squeeze(a, 2)")
c2 = torch.squeeze(a, 2)
print(c2)

output

torch.squeeze(a, 0)
tensor([[[ 1],
         [ 3],
         [ 5]],

        [[ 1],
         [ 2],
         [ 3]],

        [[ 5],
         [ 1],
         [ 3]],

        [[ 4],
         [ 2],
         [ 7]],

        [[98],
         [11],
         [ 2]]])

torch.squeeze(a, 1)
tensor([[[ 1],
         [ 3],
         [ 5]],

        [[ 1],
         [ 2],
         [ 3]],

        [[ 5],
         [ 1],
         [ 3]],

        [[ 4],
         [ 2],
         [ 7]],

        [[98],
         [11],
         [ 2]]])

torch.squeeze(a, 2)
tensor([[ 1,  3,  5],
        [ 1,  2,  3],
        [ 5,  1,  3],
        [ 4,  2,  7],
        [98, 11,  2]])

まとめ

次元拡張
np.expand_dims -> torch.unsqueeze

次元縮小
np.squeeze -> torch.squeeze