TensorFlow DatasetsでImageNet2012を力任せに読み込む話

この記事について

かなり限界なネット環境でImageNet 2012データセットをTensorFlow Datasetsのダウンロード機能に頼らず用意した時の記録です。
付属のダウンロード機能はレジューム機能無しで150GB近いデータをダウンロードするため、私の環境では完遂出来そうに無かったため、この記事のような操作を行いました。
以下の方には少し役に立つかもしれません。

• tfds.load('imagenet2012')のダウンロード速度が遅い
• 途中で何度もダウンロードに失敗した
• だけど何とかImageNet 2012データセットを使ってみたい
• そのうえでTensorFlow Datasetsで簡単に読み込みたい

TensorFlow Datasetsとは

TensorFlowにはTensorFlow Datasetsという便利なデータセット読み込み用のモジュールがあります。
これを使えば、リポジトリに予め用意されているデータセットを超簡単に読み込む事が出来ます。

MNISTを読み込む

import tensorflow_datasets as tfds
datasets = tfds.load('mnist')
for image, label in datasets['train']:
    # たったこれだけで読み込める
    pass

ImageNetを読み込もうと思った話

MNISTやCIFAR10だけでなく、もっと画像サイズの大きなデータセットに手を出してみたくなったので、リポジトリに用意されていて知名度のあるImageNet 2012(ILSVRC2012)データセットを読み込もうとした時でした。
前述の通り、読み込みはたった1行書くだけ、自動ダウンロードしてくれるので待ってれば使えるようになるはずです。……はずでした。

ダウンロードが上手くいかない

私の回線では推定残り時間が20時間~100時間の間をウロウロしていて時間が掛かりそうだったので気長に待つ事に。
しかし、1回目のダウンロードは停電で中断、2回目は3割程度の所で何故かダウンロードが停止、3回目は1割ちょいの所で停止、というように何度もダウンロードに失敗しました。
しかも、TensorFlow Datasetsのダウンロード機能にレジューム機能は無いため、毎度最初から。

よく考えると、150GB近いファイルをレジューム機能無しでダウンロードするのはかなり無謀な気がしてきたので、__ダウンロード部分は手動で行って、なんとかそれをTensorFlow Datasetsに認識させよう__という作戦に出ました。

実行環境

一般的なWindows10の環境です。

• Windows10 Pro 1809
• NVIDIA GeForce GTX 1060 6GB
• Python 3.7.4
• TensorFlow GPU 2.0.0 rc0
• TensorFlow Datasets 1.2.0

解決策

手順としては、手動ダウンロード、データセットロード用のスクリプトの改変、データセットのセットアップ、スクリプトの改変を戻す、という流れです。
途中の操作はあまりエレガントではない気がしますが、元に戻すため無かった事になるので大丈夫です。

データセットの手動ダウンロード

まずは適当な場所にILSVRC2012_img_train.tarとILSVRC2012_img_val.tarをダウンロードします。
基本的に150GB近いデータをダウンロードするので、レジューム機能が有るダウンローダを利用するのが望ましいです。

私は研究者向けのデータセット共有促進のためのサイト「Academic Torrents - ImageNet 2012」を利用しました。
Torrentの利用方法については各自で調べてください。

ダウンロードしたファイルはC:\Downloadsに保存する事にします。

Downloads

Cドライブ
| - Downloads
| | - ILSVRC2012_img_train.tar
| | - ILSVRC2012_img_val.tar

ロード用のスクリプトの改変

ImageNet 2012をロードするスクリプトは\(Pythonインストールディレクトリ)\Lib\site-packages\tensorflow_datasets\image\imagenet.pyにあります。
Pythonのインストール先は各自で読み替えて、適当なエディタで開いてください。

ファイルを開くと、_split_generators()の定義が有るはずので、探してください。
この関数がダウンロードからファイルの設置を行います。
そして、この関数の先頭を以下のように変更する事で、ダウンロードしてくるはずのtrain_pathとval_pathを直接指定します。

変更前-(Pythonインストールディレクトリ)\Lib\site-packages\tensorflow_datasets\image\imagenet.py

def _split_generators(self, dl_manager):
    train_path, val_path = dl_manager.download([
        '%s/ILSVRC2012_img_train.tar' % _URL_PREFIX,
        '%s/ILSVRC2012_img_val.tar' % _URL_PREFIX,
    ])
    if not tf.io.gfile.exists(train_path) or not tf.io.gfile.exists(val_path):
      ...

変更前-(Pythonインストールディレクトリ)\Lib\site-packages\tensorflow_datasets\image\imagenet.py

def _split_generators(self, dl_manager):
    # train_path, val_path = dl_manager.download([
    #     '%s/ILSVRC2012_img_train.tar' % _URL_PREFIX,
    #     '%s/ILSVRC2012_img_val.tar' % _URL_PREFIX,
    # ])
    # TODO: 元の代入をコメントアウトして手動ダウンロードしたファイルのパスを直接指定する
    train_path = 'C:\\Downloads\\ILSVRC2012_img_train.tar'
    val_path = 'C:\\Downloads\\ILSVRC2012_img_val.tar'
    if not tf.io.gfile.exists(train_path) or not tf.io.gfile.exists(val_path):
      ...

変更し終えたら保存してください。

tfds.load('imagenet2012')の実行

適当なファイルにimportしてloadするスクリプトを書いて実行してください。
データセット用のディレクトリは各自の使用している物を使用してください。
ファイルを展開するために**143GB(153,745,173,498B)**必要になるので、空き容量を事前に確認してください。
ここではC:\Datasetsとします。

main.py

import tensorflow_datasets as tfds
datasets = tfds.load('imagenet2012', data_dir='C:\\Datasets')

実行

$ python main.py

実行すると、ダウンロードの部分をスキップしてファイルの展開と設置が始まるはずです。

ロード用スクリプトを元に戻す

先ほど改変したimagenet.pyを変更前に戻してください。

完了

以上の手順を踏むと、通常の手順でデータセットを設置した時と同じ状態になります。
後は、通常通りtfds.load('imagenet2012')で読み込んで使うだけ。
お疲れ様でした。

参考文献

TensorFlow Datasets
github.com - tensorflow_datasets/image/imagenet.py
Academic Torrents
27TB以上の研究用データセットをBitTorrentで共有する「Academic Torrents」