DeepNudeの学習済みファイルをBarracudaで使用するためにonnxファイルに変換する

環境

python3.7

概要

Unityの機械学習ライブラリであるBarracudaでDeepNudeの学習済みデータを使用するためにonnxファイルに変換します。

以下の関数をgan.py内にあるDeepModelクラスに追加します。

	def export_onnx(self):
		self.netG.eval()
		
		batchSize = self.opt.batchSize
		imageChannel = self.opt.input_nc
		width = 512
		height = 512
		input = torch.randn(batchSize, imageChannel, width, height)
		checkpointPath = self.opt.checkpoints_dir
		basename_without_ext = os.path.splitext(os.path.basename(checkpointPath))[0]
		print(basename_without_ext)
		outputName = basename_without_ext + ".onnx"
		print(outputName)
		torch.onnx.export(
				self.netG,							# モデル
				input,								# モデルの入力
				outputName,						# ONNXファイル名
				export_params=True,			# 重みをONNXファイルに保存
				opset_version=9,				# ONNXのバージョン
				do_constant_folding=True,	# 最適化のために定数の折りたたみの実行
				input_names = ['X'],			# モデルの入力名
				output_names = ['Y'],		# モデルの出力名
				verbose=True						#変換中の詳細ログ
		)

run.py内の model.initialize(opt)の後にmodel.export_onnx()を挿入して呼び出すとファイルに保存されます。

その学習済みデータ（水着部分を緑色にする）を使用してUnityで水玉コラージュを作成してみました。
詳しくはリンクを参照してください。

参照

https://github.com/zhengyima/DeepNude_NoWatermark_withModel
https://kenken5050.hatenablog.com/entry/2022/05/28/111144