ImportError: No module named apex.fp16_utilsの解決方法（備忘録）

上記のエラーについて、今回は備忘録的に残しておきます！

1. 背景
2. 過程（ソースコードを見てみる）
3. 結論

背景

Google Colabで学習済みのSSDを推論させるべく、必要なコードをColabに移植したところ、上記のエラーImportError: No module named apex.fp16_utilsが出てきました🫠
「これ、どうしたら良いんや、、、ライブラリのエラーとかどうしようもないやん、、、」
色々とIssueは立っていてインストール方法を変えるだのなんだのあったので全部試しましたがダメでした、、、

といったところで、apex（Github）のソースコードを見ることにした次第です。

過程（ソースコードを見てみる）

今回、importできないのはapexの中のfp16_utilsです。つまり、ライブラリであるかどうかは置いといて、Pythonの構造的にはapex配下のfp16_utilsというディレクトリあるいは.pyファイルを確認すれば良いわけです。

実際にNVIDIA/apexを見てみると、以下のようになってました。
案の定apex/fp16_utilsの構造ですね。

更にfp16_utilsの中を見て、使用されているコードを引っ張って実質ライブラリを導入せずに使えるかどうか見てみます。（PyTorchはColabにデフォルトで入っているので考慮しなくてOK）

大体モジュールをExportしているのは__init__.pyなので
そこから今回欲しいコードを入手します。

SSDで使いたかったのはnetwork_to_half関数なので、そちらをメインで探します。

__init__.pyを見ると、Githubの便利機能でapex/fp16_utils/fp16_util.pyにあるとのこと。見てみましょう！

ここにありますね。よく見ると、依存している部分（Import）は全てPyTorchのみで
他のライブラリやApex内の他ファイルからImportしていないので、PyTorchが使える環境で、この関数に関するコードをそのままコピペすれば動作しますね！

結論

必要だった以下のコードをapexから取得し、そのままコピペ。

import torch
import torch.nn as nn

class tofp16(nn.Module):
    """
    Utility module that implements::

        def forward(self, input):
            return input.half()
    """

    def __init__(self):
        super(tofp16, self).__init__()

    def forward(self, input):
        return input.half()

def BN_convert_float(module):
    """
    Utility function for network_to_half().

    Retained for legacy purposes.
    """
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm) and module.affine is True:
        module.float()
    for child in module.children():
        BN_convert_float(child)
    return module

def network_to_half(network):
    """
    Convert model to half precision in a batchnorm-safe way.

    Retained for legacy purposes. It is recommended to use FP16Model.
    """
    return nn.Sequential(tofp16(), BN_convert_float(network.half()))