Atcoder お誕生日コンテストX - 「この問題はほんとうにひどい問題であるため，できれば先に他の問題のほうをお楽しみいただければと思っておりまして，ですので他の問題を通し終えて暇になり，かつその暇を」をCNNで解く

問題

atcoder.jp

問題概要

括弧を含む四則演算の式が与えられるので、答えを整数で出力する。ただし、式は画像で与えられる。

式の画像は、Courier Primeというフォントの各文字に対して回転、縦横の拡大（縮小）、せん断変形、平行移動を行って生成される。さらに、確率 $P=0.05$ で画素を反転させるノイズが生成される。

この問題について

2014年2月に開催されたJAPLJさんのお誕生日コンテストで出題された問題。

AtCoderで開催されるお誕生日コンテストではいわゆるギャグ寄りの問題が出題される傾向にある（というよりこのコンテストが原点？）が、この問題はその中でもかなり強いギャグである。

僕がこの問題を知ったのは競プロを始めてしばらく経った2018年ごろだったが、かなり印象に残っている問題だった（「解けるわけないのになぜかACしている人がいる」と思っていた）。

最近ニューラルネットワークをC++に埋め込むことについて考えていたときにこの問題のことを思い出し、チャレンジしてみることにした。

制約が当時と異なる点について

問題の制約は問題文中に書かれているが、そのなかに「AtCoder で送信できるソースコードは UTF-8 で 60,000 文字以下なので，文字数制限に注意してプログラムを書くこと．」とある。しかし、現在のAtCoderにおけるコード長制限は512KiBであり、これは1バイト文字なら500000文字以上に相当する。今回は、結果的に当時の制約のなかで解くことができた。

解法

問題設定

この問題は「画像データから数字と記号の情報を取り出す」「得られた式をパースして計算する」という2段階に分けられる。後者は競プロでは構文解析の超典型問題なので難しい部分は前者だが、前者は前者で機械学習における典型データセットであるMNISTと非常によく似ている。

記号が6種類あるのでMNISTよりクラスの数が少し増えているのと、満点である500点を取るためにはモデルの正答率100%が必要である。とはいえMNISTにはどう考えてもラベル付けをミスっているとしか思えないデータや既存の数字のどれにも似ていないデータなどが含まれているため、それに比べると正答率100%は目指しやすいと考えられる。

学習データセットの生成

学習はPythonで行うので、Pillowを使ってデータセットを生成した。

実際にクエリとして与えられる画像と同じように、Courier Primeの各文字に対してランダムに回転、拡縮、せん断変形、平行移動を行っていく（回転角は少し多めに±20°まで取ってみた）。推論の際にはクエリ画像にメジアンフィルタを適用してノイズを除去するが、その際に文字の形が変わりがちなので、あらかじめ $P=0.05$ でノイズを生成してからメジアンフィルタを適用したものを教師データとして用いた。実際に与えられるクエリ画像では文字が途切れているということはないが、~~面倒なので~~強いモデルを作るために修正は特に行わなかった。データセットは各文字5000枚ずつ生成した。

モデルの作成

学習にはPyTorchを用いた。モデルは以下のように定義した。

class Model(nn.Module):
    def forward(self, x):
        h = F.relu(self.conv1(x))
        h = F.relu(self.conv2(h))
        h = F.relu(self.conv3(h))
        h = self.pool(h)
        h = h.view(-1, 16*6*3)
        h = self.linear1(h)
        o = self.linear2(h)
        return o

    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=4, kernel_size=3, stride=2).to(device)
        self.conv2 = nn.Conv2d(in_channels=4, out_channels=8, kernel_size=3, stride=2).to(device)
        self.conv3 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1).to(device)
        self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)).to(device)
        self.linear1 = nn.Linear(in_features=16*6*3, out_features=20).to(device)
        self.linear2 = nn.Linear(in_features=20, out_features=16).to(device)

生成したデータセット80000枚をtrainデータ64000枚とtestデータ16000枚に分割して学習率0.001で200epochほど学習させると、trainデータに対する正答数が64000、testデータに対する正答数が15983となった。accuracyは99.8%だが、データセットはクエリ画像よりも少し難しめのものを生成していることを考えるとまあまあな気がした。学習時間はデスクトップPC（with RTX3080）で10分くらいだった気がする（ちゃんと計ってない）

モデルをC++に移植する

AtCoderではPyTorchは使えないので、推論を行うには別の方法を用いる必要がある（追記：今度の言語アップデートで入るらしい？）。今回は当初のモチベーション通りC++にモデルを埋め込んでみる。

Python側で学習したモデルをエクスポートするとき、パラメータをそのままリスト形式で書き出すという方法もなくはないが、バイナリにパラメータを詰め込むほうが圧倒的に効率がよい。

PyTorchのパラメータはデフォルトでは4バイトの浮動小数点数なので、10000パラメータを埋め込んだバイナリは約40000バイトになる。今回はインポートも自分で書くのでパラメータを自分で決めた順番で並べるだけでよいが、適当な場所にパラメータ数の情報などを整数で挟み込んでおくと読み込み時にずれていないかどうかなどが判断できて便利。

通常の環境であればこのバイナリをファイルとして保存してC++側で読み込むという処理で問題ないが、AtCoderではファイル一枚しか提出できないためモデルを文字列として埋め込む必要がある。今回はバイナリを文字列に変換するためにbase64によるエンコードを用いた。base64では6バイトを1文字で表すため、結局10000パラメータは54000文字程度になる。

    def get_linear_params(self, linear):
        param = struct.pack("<I", linear.weight.nelement())
        for val in linear.weight.data.flatten().tolist():
            param += struct.pack("<f", val)
        param += struct.pack("<I", linear.bias.nelement())
        for val in linear.bias.data.flatten().tolist():
            param += struct.pack("<f", val)
        return param

    def get_conv2d_params(self, conv2d):
        param = struct.pack("<I", conv2d.weight.nelement())
        for val in conv2d.weight.data.flatten().tolist():
            param += struct.pack("<f", val)
        param += struct.pack("<I", conv2d.bias.nelement())
        for val in conv2d.bias.data.flatten().tolist():
            param += struct.pack("<f", val)
        return param

    def get_parameter_string(self):
        # C++モデル用のbase64を出力する
        params = bytes()

        params += self.get_conv2d_params(self.conv1)
        params += self.get_conv2d_params(self.conv2)
        params += self.get_conv2d_params(self.conv3)
        params += self.get_linear_params(self.linear1)
        params += self.get_linear_params(self.linear2)

        return base64.b64encode(params).decode()

次に、C++側でモデルを用意する。 AtCoderで使えるC++の機械学習ライブラリは多分ないので、推論部分を自前で実装していく。線形層を実装するのは比較的簡単だが、畳み込み層の処理を実装するのはかなり面倒くさい。とはいえ、実際にPyTorchで動かした推論処理と比較しながら実装していけばなんとかなった。

C++は適当に書いても充分高速なので、今回のケースであれば愚直な実装でも実行時間にはかなり余裕があった。さらに、モデルのパラメータをインポートするためにbase64のデコードとバイナリの読み込みを追加で実装した。

これで、 $38\times65$ の画像を読み込んで描かれている記号を推論するモデルが用意できた。

残りの部分を書く

本質の部分が実装できたので、あとは残りの部分を書いていく。

ジャッジから与えられるクエリ画像にはノイズがかかっているので、最初に $3\times3$ のメジアンフィルタをかけてノイズを除去する。

作成したモデルは $38\times65$ の画像しか推論できないので、文字の場所を特定して画像を切り出す必要がある。文字は基本的に大きな連結成分になっているので、黒いピクセルがある列が5列以上並んでいれば文字であるとして判断した。各画像を推論すれば文字列の式が得られるので、あとは構文解析を書けば完成。

これを提出してもWAが1～2ケース出る場合があるが、何度か学習をやり直してモデルガチャをすれば数回でACできた（60000文字以内に収めようとした場合。現在のコード長制限をフルで使って大きなモデルを使った場合は一発でACできた。）

提出

atcoder.jp

感想

軽い気持ちで手を付けたら2日くらいかかってしまった（大半はCNNの実装）

昔この問題を競プロとして見たときはタイトルの通り「ほんとうにひどい問題」だと思ったが、ニューラルネット埋め込みの課題として見れば教育的な良問だった。

なんにせよ、昔から気になっていた問題をACできてよかった。