TechFULの中の人

TechFULスタッフ・エンジニアによる技術ブログ。IT関連のことやTechFULコーディングバトルの超難問の深掘り・解説などを毎週更新

PyTorchを使ったワンライナー機械学習


こんにちは、TechFULでアルバイトをしている、berryberryです。 現在は、TechFUL PROと呼ばれる人工知能などの最先端技術を測定/学習するためのサービスに関連した問題の作成や英訳作業をおこなっています。

私は大学院で深層学習に関係する研究をおこなっているのですが、特にPyTorchをよく利用しています。 PyTorchは深層学習をおこなう上で非常に便利なPythonライブラリですが、ソースコードが少し長くなってしまうという問題があります。 ちょっと思いついたアイデアを試すのに長いソースコードを書くのは面倒ですよね?なんならエディタを開くのも面倒だと思います。(そんなことはない?)

今回の記事では、PyTorchをシェルからワンライナーで実行することを試みます。シェルさえあれば、エディタなんていらんのです。 そのため、PyTorchの基本的なところは今回の記事では紹介しません。

TechFUL PROではPyTorchの基礎を学べるラーニング問題を公開していますので、PyTorchの基礎が分からない人はぜひともラーニング問題を解いてみてください。

基本のソースコード

まずは、PyTorchを使ってcos関数を学習するプログラムを見てみましょう。 このプログラムをシェルからワンライナーで叩けるように書き換えるのが目標となります。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import data

# seedを固定
torch.manual_seed(123)

# 学習データの作成
x = torch.linspace(-1, 1, 1000, dtype=torch.float).reshape(-1, 1)
y = torch.cos(x) + 0.01*torch.randn_like(x)

# データを訓練データと検証用のデータに分割
train_data, val_data = data.random_split(data.TensorDataset(x, y), [900, 100])

# それぞれのデータをDataLoaderに渡す
# バッチサイズは32に設定
train_loader = data.DataLoader(train_data, 32, shuffle=True, drop_last=True)
val_loader = data.DataLoader(val_data, 100)

# モデルの定義
model = nn.Sequential(
    nn.Linear(1, 128),
    nn.ReLU(),
    nn.Linear(128, 128),
    nn.ReLU(),
    nn.Linear(128, 1),
    nn.Tanh(),
)

# オプティマイザを設定
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.001)

# 20 epochだけ学習
for epoch in range(20):
    train_losses = []
    model.train()
    # 訓練データを1周
    for input, target in iter(train_loader):
        # 勾配をリセット
        optimizer.zero_grad()

        # ロスの計算
        output = model(input)
        loss = F.mse_loss(output, target)

        # ロスに基づいてモデルを更新
        loss.backward()
        optimizer.step()

        train_losses.append(loss.detach().cpu())

    val_losses = []
    model.eval()
    # 検証データを1周
    for input, target in val_loader:
        output = model(input)
        loss = F.mse_loss(output, target)
        val_losses.append(loss.detach().cpu())

    # 訓練データと検証データそれぞれに対するロスの平均を計算し、出力。
    train_loss = torch.stack(train_losses).mean().item()
    val_loss = torch.stack(val_losses).mean().item()

    print(f"Epoch: {epoch+1} | Train Loss: {train_loss:.2f} | Val Loss: {val_loss:.2f}")

このコードを動かすと、以下のような結果が得られました。 検証データに対するロス(Val Loss)が0になっているので、うまく学習できていそうです。

Epoch: 1 | Train Loss: 0.06 | Val Loss: 0.04
Epoch: 2 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 3 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 4 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 5 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 6 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 7 | Train Loss: 0.03 | Val Loss: 0.02
Epoch: 8 | Train Loss: 0.01 | Val Loss: 0.01
Epoch: 9 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 10 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 11 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 12 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 13 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 14 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 15 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 16 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 17 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 18 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 19 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 20 | Train Loss: 0.00 | Val Loss: 0.00

ワンライナーで動かすためのテクニック

次に、上記のプログラムをワンライナーで書くために有効なテクニックを紹介します。 上のプログラムを書き換えるうえで問題となりそうなのは、

  • モデルなどの定義
  • 訓練のforループ
  • ライブラリのインポート

あたりでしょうか?これらを解決するために、本記事ではセイウチ演算子とリスト内包表記を利用します。

セイウチ演算子

セイウチ演算子はPython3.8から導入された演算子で、変数への代入と実行を同時に利用することができます。 たとえば、リストの要素数が3である場合にリストの要素数を表示するプログラムは以下のように書けます。

x = [1,2,3]
if len(x) == 3:
    print(len(x))

これを、セイウチ演算子を使って書き換えると以下のようになります。

x = [1,2,3]
if (n:=len(x)) == 3:
    print(n)

このセイウチ演算子を使うと、同じリスト内でモデルを定義し、更新することができるようになります。 具体例として、一番シンプルなモデルの定義、更新をおこなうプログラムを考えます。 これは、以下のように書けます。

model = nn.Sequential(...)
optimizer = torch.optim.Adam(...)

loss = F.mse_loss(model(input), target)
loss.backward()
optimizer.step()

このプログラムをセイウチ演算子で書き換えると、以下のように書けます。 見やすくするために改行を適宜入れているのですが、実は下のプログラムは改行部分を削除して1行にしても問題なく動作します。

[
    model := nn.Sequential(...),
    optimizer := torch.optim.Adam(...),
    loss := F.mse_loss(model(input), target),
    loss.backward(),
    optimizer.step(),
]

リスト内包表記

Pythonユーザにとって、リスト内包表記は身近なものかもしれません。 リスト内包表記とは、for文を1行で書くことができる表記方法です。

例えば、1~10までの整数を格納したリストを作るとします。 この時、for文を使うと以下のようになります。

x = []
for i in range(10):
    x.append(i+1)

これをリスト内包表記で書くと、以下のようになります。

x = [i+1 for i in range(10)]

リスト内包表記は、ソースコードをスッキリと書けるだけでなく、appendを呼ぶ必要がないため実行時間が短くなるというメリットもあります。 PyTorchを使ったソースコードではfor文がよく出てくるため、ワンライナー機械学習をするときにはリスト内包表記が有効になります。

ライブラリのインポート

ライブラリのインポートをワンライナーで書くには、__import__とセイウチ演算子を組み合わせることで解決できます。 以下の2つのプログラムはどちらも同じように動きます。

import torch
print(torch.rand(size=(1,)))
[
    torch = __import__("torch"),
    print(torch.rand(size=(1,)))
]

ワンライナーへの書き換え

それでは、学習のコードをワンライナーに書き換えていきましょう。基本は、セイウチ演算子を使ってモデルなどを定義することと、リスト内包表記で学習のループを書くことです。

...できました!一番始めのプログラムと見比べてもらうと、対応関係もわかりやすいと思います。見やすくするために改行をいれていますが、改行を除いても問題なく動きます。

[
    torch := __import__("torch"),
    nn := torch.nn,
    F := torch.nn.functional,
    data := torch.utils.data,
    
    x := torch.linspace(-1, 1, 1000, dtype=torch.float),
    y := torch.cos(x) + 0.01*torch.randn_like(x),
    datasets := data.random_split(data.TensorDataset(x, y), [900, 100])
    train_loader := data.DataLoader(
        datasets[0], 32, shuffle=True, drop_last=True),
    val_loader := data.DataLoader(datasets[1], 100),
    model := nn.Sequential(
        nn.Linear(1, 128),
        nn.ReLU(inplace=True),
        nn.Linear(128, 128),
        nn.ReLU(inplace=True),
        nn.Linear(128, 1),
        nn.Tanh(),
    ),
    optimizer := torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.001),
    [
        [
            train_losses := [],
            val_losses := [],

            [
                [
                    model.train(),
                    optimizer.zero_grad(),
                    loss := F.mse_loss(model(input), target),
                    loss.backward(),
                    optimizer.step(),
                    train_losses.append(loss.detach()),
                ] for input, target in iter(train_loader)
            ],

            [
                [
                    model.eval(),
                    val_losses.append(F.mse_loss(model(input), target).detach()),
                ] for input, target in iter(val_loader)
            ],

            print("Epoch: {} | Train Loss: {:.2f} | Val Loss: {:.2f}".format(
                epoch+1,
                torch.stack(train_losses).mean().item(),
                torch.stack(val_losses).mean().item(),
            ))

        ] for epoch in range(20)
    ]
]

あとは上記のプログラムからコメント部分と改行を除いた以下のコマンドをシェルでたたけば、機械学習ができます。 これで、毎回長いプログラムを書かずともシェルからワンライナー機械学習できるようになりました。 TechFUL PROでは、機械学習を使う問題がいくつか公開されていますので、ぜひともワンライナー機械学習を試してみてください!

python -c '[torch := __import__("torch"), nn := torch.nn, F := torch.nn.functional, data := torch.utils.data, torch.manual_seed(123), x := torch.linspace(-1, 1, 1000, dtype=torch.float).reshape(-1, 1), y := torch.cos(x) + 0.01*torch.randn_like(x), datasets := data.random_split(data.TensorDataset(x, y), [900, 100]), train_loader := data.DataLoader(datasets[0], 32, shuffle=True, drop_last=True), val_loader := data.DataLoader(datasets[1], 32), model := nn.Sequential(nn.Linear(1, 128), nn.ReLU(inplace=True), nn.Linear(128, 128), nn.ReLU(inplace=True), nn.Linear(128, 1), nn.Tanh()), optimizer := torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.001), [[train_losses := [], val_losses := [], [[model.train(), optimizer.zero_grad(), loss := F.mse_loss(model(input), target), loss.backward(), optimizer.step(), train_losses.append(loss.detach())] for input, target in iter(train_loader)], [[model.eval(), val_losses.append(F.mse_loss(model(input), target).detach())] for input, target in iter(val_loader)], print("Epoch: {} | Train Loss: {:.2f} | Val Loss: {:.2f}".format(epoch+1, torch.stack(train_losses).mean().item(), torch.stack(val_losses).mean().item()))] for epoch in range(20)]]'

結果

Epoch: 1 | Train Loss: 0.06 | Val Loss: 0.04
Epoch: 2 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 3 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 4 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 5 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 6 | Train Loss: 0.04 | Val Loss: 0.04
Epoch: 7 | Train Loss: 0.03 | Val Loss: 0.02
Epoch: 8 | Train Loss: 0.01 | Val Loss: 0.01
Epoch: 9 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 10 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 11 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 12 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 13 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 14 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 15 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 16 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 17 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 18 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 19 | Train Loss: 0.00 | Val Loss: 0.00
Epoch: 20 | Train Loss: 0.00 | Val Loss: 0.00

最後に

反省点としては、リスト内包表記を有効活用できていないことです。実はリスト内包表記はチューリング完全を満たしており、リスト内包表記だけを使ってモデルを学習することもできるはずです。 しかし、今回は訓練のループ以外ではリスト内包表記を使えていません... まあ、ワンライナーで動くプログラムではあるので良しとしましょう。

最後に、ワンライナー機械学習する意味は特に無いのでちゃんとプログラムを書いて動かすのがよいと思います。 間違っても、研究や業務などでワンライナーにするのはやめましょう。