よっしーの私的空間

機械学習を中心に興味のあることについて更新します

TensorFlowとPytorchにおける画像のデータ構造の違いと変換方法について

本記事では以下についてまとめます。

  1. TensorFlowとPytorchにおける画像のデータ構造の違いについて解説
  2. データ構造の相互に変換する方法について解説

1.TensorFlowとPytorchの画像のデータ構造の違い

1.1.概要

画像データは4次元配列で表現され、各次元の要素で以下を表現します。

  • バッチ数(画像の枚数)
  • 高さ
  • チャネル(色の成分の数(例:カラー画像はRGBでチャネル数3))

TensorFlowとPytorchで上記の要素の順番が異なります。

  • TensorFlowの場合
     [バッチ数, 高さ, 幅 , チャネル]
  • Pytorchの場合
     [バッチ数, チャネル, 高さ, 幅]

TensorFlow用の画像のデータ構造は「channel last」と呼び、Pytorch用は「channel first」と呼びます。

1.2.具体例を使ってデータ構造の違いを確認

TensorFlow用の画像

配列の中が見えやすいように、numpyを使って、Tensorflow用の画像を作成してみます。

# 高さ2×幅6のカラー画像(チャネル数3)の4枚を作成
image_channel_last = np.array([
    # 1個目の画像
    [[[0,255,0],[0,0,255],[255,0,0],[0,255,0],[0,0,255],[255,0,0]],
     [[255,0,0],[0,255,0],[0,0,255],[255,0,0],[0,255,0],[0,0,255]]], 
    # 2個目の画像
    [[[0,255,255],[255,0,255],[255,255,0],[0,255,255],[255,0,255],[255,255,0]],
     [[255,0,255],[255,255,0],[0,255,255],[255,0,255],[255,255,0],[0,255,255]]],
    # 3個目の画像
    [[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]],
     [[255,255,255],[255,255,255],[255,255,255],[255,255,255],[255,255,255],[255,255,255]]],
    # 4個目の画像
    [[[255,255,255],[255,255,255],[255,255,255],[255,255,255],[255,255,255],[255,255,255]],
     [[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]]],
])

上記をnp.shape関数で各次元の要素数を確認してみると以下の通りとなり、channel lastになっていることが分かります。
f:id:t-yoshi-book:20210607174400p:plain

Pytorch用の画像

続いてPytorch用の画像の配列を作成してみます。

# 高さ2×幅6のカラー画像(チャネル数3)の4枚を作成
image_channel_first = np.array([
    # 1個目の画像
    [[[  0,   0, 255,   0,   0, 255],[255,   0,   0, 255,   0,   0]], #赤成分
     [[255,   0,   0, 255,   0,   0],[  0, 255,   0,   0, 255,   0]], #緑成分
     [[  0, 255,   0,   0, 255,   0],[  0,   0, 255,   0,   0, 255]]],#青成分
    # 2個目の画像
    [[[  0, 255, 255,   0, 255, 255],[255, 255,   0, 255, 255,   0]],
     [[255,   0, 255, 255,   0, 255],[  0, 255, 255,   0, 255, 255]],
     [[255, 255,   0, 255, 255,   0],[255,   0, 255, 255,   0, 255]]],
    # 3個目の画像
    [[[  0,   0,   0,   0,   0,   0],[255, 255, 255, 255, 255, 255]],
     [[  0,   0,   0,   0,   0,   0],[255, 255, 255, 255, 255, 255]],
     [[  0,   0,   0,   0,   0,   0],[255, 255, 255, 255, 255, 255]]],
    # 4個目の画像
    [[[255, 255, 255, 255, 255, 255],[  0,   0,   0,   0,   0,   0]],
     [[255, 255, 255, 255, 255, 255],[  0,   0,   0,   0,   0,   0]],
     [[255, 255, 255, 255, 255, 255],[  0,   0,   0,   0,   0,   0]]]
])

こちらもnp.shapeで各次元の要素数を確認してみます。channel firstになっていることが分かります。
f:id:t-yoshi-book:20210607174700p:plain

2.データ構造の変換方法

np.transpose関数を使用することで非常に簡単に変換が可能です。

2.1.Channel last ⇒ Channel first

image_channel_first = np.transpose(image_channel_last , [0,3,1,2])

2.2.Channel first ⇒ Channel last

image_channel_last = np.transpose(image_channel_first , [0,2,3,1])