ColabでのUnavailable Errorの対処法

はじめに

ColabのTPUインスタンス上でtf.kerasを用いてCNNを学習させていたのですが「Unavailable Error : OS Error」や「Unavailable Error : Socket closed」などが発生して困っていました。やっかいなことにこれらのエラーは再現性がなく学習を始めてすぐに発生することもあれば一切発生しないこともあります。

環境

Google Colaboratory(tpuインスタンス)
tensorflow 1.13.1

原因

はっきりとした原因は不明ですがどうやらGoogleDriveとのIO周りの問題のようです。
私は上記のCNNを学習させる際にデータをGoogleDrive上に置いておき、下記のコードのようにColab上にマウントし、flow_from_directoryを用いてfit_generatorによって学習を行っていました。

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from google.colab import drive
from oauth2client.client import GoogleCredentials

drive.mount('/content/drive')
savedir ="drive/データの保存先"
.
.
.
# GoogleDriveから都度都度読み込む
train_datagen = ImageDataGenerator(rescale = 1.0 / 255)
validation_datagen = ImageDataGenerator(rescale = 1.0/ 255)

train_generator = train_datagen.flow_from_directory(
                  savedir+"/train",
                  target_size = (256, 256),
                  batch_size=batch_size,
                  class_mode = "categorical")

validation_generator = train_datagen.flow_from_directory(
                  savedir+"/validation",
                  target_size = (256, 256),
                  batch_size=batch_size,
                  class_mode = "categorical")
.
.
.
history = model.fit_generator(train_generator,
                             steps_per_epoch = train_datasize / batch_size,
                             epochs=100,
                             verbose=1,
                             validation_data=validation_generator,
                             validation_steps = valid_datasize / batch_size,
                             shuffle=True)

しかしこのやり方では学習を進める過程でその都度GoogleDriveから必要なデータを読み込んでくるみたいです。ColabとGoogleDriveの間のIOは遅い上にどうも不安定で時折失敗してしまいます。(一度読み込んできたデータはColab上にキャッシュされるようなので2Epoch目からは普通に進んでいきます。)
どうもこのIOの不安定さがUnavailable Errorの原因のようです。

対処法

GoogleDriveとの通信を最小限に抑えます。
学習データや評価データを事前にtarファイル等にまとめておき学習を始める前にそれらのファイルだけをGoogleDriveから持ってきてColabのローカルにデータ本体を展開しておきます。こうすることで最初にアーカイブファイルを持ってくるところだけ成功すれば、後はColab上だけで話が進んでいくのでUnavailable Errorも避けることができますし1Epoch目の学習が早くなります。

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from google.colab import drive
from oauth2client.client import GoogleCredentials

drive.mount('/content/drive')
savedir ="drive/データの保存先"
.
.
.
# tarファイルだけGoogleDriveから持ってきて必要なデータをColab上に展開する。
train_datagen = ImageDataGenerator(rescale = 1.0 / 255)
validation_datagen = ImageDataGenerator(rescale = 1.0/ 255)
import tarfile
train = tarfile.open(savedir+"train.tar", mode = 'r')
train.extractall()
validation = tarfile.open(savedir+"validation.tar", mode = 'r')
validation.extractall()
.
.
.
# Colab上のデータを読み込む
train_datagen = ImageDataGenerator(rescale = 1.0 / 255)
validation_datagen = ImageDataGenerator(rescale = 1.0/ 255)

train_generator = train_datagen.flow_from_directory(
                  "train",
                  target_size = (256, 256),
                  batch_size=batch_size,
                  class_mode = "categorical")

validation_generator = train_datagen.flow_from_directory(
                  "validation",
                  target_size = (256, 256),
                  batch_size=batch_size,
                  class_mode = "categorical")
.
.
.
history = model.fit_generator(train_generator,
                             steps_per_epoch = train_datasize / batch_size,
                             epochs=100,
                             verbose=1,
                             validation_data=validation_generator,
                             validation_steps = valid_datasize / batch_size,
                             shuffle=True)

まとめ

Unavailable Errorに関しては調べてもこれという情報がなく困っていました。
上の方でGoogleDriveとのIOが原因だと述べていますが、たまたま上記の対処法を試してみたところ治ったのでおそらくそういうことだろうという仮説でしかないのでもしも実際のところをご存知の方がおられましたら教えてください。m(_ _)m
同じ症状で悩んでおられる方はぜひ一度お試しあれあれ。

追記(愚痴)

体感的にオートエンコーダ用途で使用すると何かと不具合が多い気がします。
謎だ...