Mojoのファーストインプット〜機械学習によるPythonとの速度比較(したかった)

はじめに

この記事は、株式会社システムアイのAdvent Calendar 2023 の 11 日目の記事となります。

皆さんは、Mojoというプログラミング言語は知っていますか？IT 業界にいる方なら聞いたことはありますよね！

2023年5月2日に、LLVM や Swift の生みの親として知られる Chris Lattner 氏が率いる Modular 社 が開発したコンパイル言語です！

その特徴は、なんといっても実行速度！Python の 68,000 倍の実行速度が出ます！アセンブリよりも早いという理論上意味の分からない数値ですが、GPU による演算やマルチコアなど、あらゆる速度最適化によって、非常に高速な処理を実現しているようです。現代技術を組み込んだ新言語の強みが出てますね！

「Python の完全上位互換」「Python とは記法も思想も全く違う」「C と C++ の関係」「JavaScript と TypeScript の関係」など、良い意味でも悪い意味でも話題になっている言語ですが、Python 大好きマンの私としては触れないわけにはいきません！（笑）

ということで、今回 Mojo を触ってみました！

環境

• Google Colaboratory (2023/12/9)
• T4 GPU
• NVIDIA-SMI 525.105.17 Driver Version: 525.105.17 CUDA Version: 12.0
• Python 3.10.12
• pip 23.1.2

GPU を使用したり OS や スペック依存の問題など、本題と関係ないところでの問題が起こらないように、Colab 上で実行していきます！（Colab の使い方は難しくないため省きますが、ランタイムタイプをデフォルトから GPU に変更しています。）

目次

• Mojo インストール
• Mojo ファーストインプット
• 機械学習による Python との速度比較
• おまけ

Mojo インストール

• Mojo の公式サイト から、「Download Now」を押します。
  

• アカウント作成を促されるので、任意の方法で登録します。

• セットアップページに飛ぶので、左の「Setup on Windows」を「Install on Linux」に変更します。
  その後、Google Colab にて、表示されているコマンドに従って、SDK の更新とインストールを行います。（Google Colab では、コマンドの最初に「!」を付ける必要があります。）

# Colab のコードセルで、このコマンドを実行（SDK更新）
!sudo apt-get update && \
sudo apt-get install modular && \
modular clean && \
modular install mojo

# Colab のコードセルで、このコマンドを実行（手順2）
!curl https://get.modular.com | sh - && \
modular auth mut_<自分のMODULAR_auth>

# Colab のコードセルで、このコマンドを実行（手順3）
!modular install mojo

• Mojo コマンドがしっかりと実行できるか確認します。
  Google Colab 上でmojoコマンドを実行する際は !export MODULAR_HOME="/root/.modular" && export PATH="/root/.modular/pkg/packages.modular.com_mojo/bin:$PATH" && をコマンドの前につけて環境変数を渡す必要があります！

!export MODULAR_HOME="/root/.modular" && \
export PATH="/root/.modular/pkg/packages.modular.com_mojo/bin:$PATH" && \
mojo --version

これにて完了です。

Mojo のファーストインプット

Python との違いとして、特に特徴的な以下の点について紹介します！

• 基礎知識
• 関数定義
• borrowed と inout と owned
• 構造体
• Pythonモジュールの利用

基礎知識

前提として、Mojo は Python との高い互換性を目指しています！そのため、Python のコードの変更をあまり必要としないコード記法になっています！（便利！！）

まず違いとして、拡張子が「.mojo」になります。(そのままの拡張子で分かりやすいですね！)

また、AOT (ahead-of-time) と JIT (just-in-time) 両コンパイル方法に対応しています。また、後述する「Pythonモジュールの利用」でも紹介しますが、部分的に Python インタプリタも使用されます。

コーディング・実行方法

Python では特に何もなく順次実行で上から書いていけばよかったですが、Mojo ではエントリーポイントとして Main 関数が必須になっています。

実行方法は、AOT の場合は「mojo build mojoのファイル名」でコンパイルして「./実行ファイル名」などで実行、JIT の場合は「mojo mojoのファイル名」で実行します。

(備考) Google Colab 上での実行方法

1. セルの上部に %%writefile main.mojo を書いて、mojoファイルとして出力する
2. 下記のコマンドで Mojo コマンドを実行する

!export MODULAR_HOME="/root/.modular" && \
export PATH="/root/.modular/pkg/packages.modular.com_mojo/bin:$PATH" && \
mojo main.mojo

変数定義

Mojo は Python と違い、変数定義の際に「var（ミュータブル）」「let（イミュータブル）」が必要です。もちろん、let に後から代入することは不可能です。

# Python
a = 1
b = 2

a += 5
b += 10

print(a)  # 6
print(b)  # 12

# Mojo
fn main():
	var a = 1
	let b = 2
	
	a += 5  # できる
	b += 10  # エラーになる
	
	print(a)
	print(b)

型の定義

Mojo の型は「Int、Float、String、Bool、List、Tuple」などに対応しています。Python の型「int、float、str、bool、list、tuple」などとは別のものとして扱われ、コンパイル用に最適化され、速度も速くなります。ただ、Pythonの型にも対応しているため、既存の Python コードをそのまま利用できます。（便利！！）

また、Python では見た目だけの何の意味もない型アノテーションが、Mojo では型宣言として使用されるようになりました！

# Python
a: int = "abc"  # ただのメモなだけで、何の効力もない
print(a)  # abc

# Mojo
fn main():
	var a: Int = "abc"  # しっかりとエラーになる
	print(a)

関数定義

Mojo では、「fn 関数名()」で関数定義を書きます。先ほど同様 Python の関数定義「def 関数名()」とは別のものとして扱われ、コンパイル用に最適化され、速度も速くなります。ただ def にも対応しているため、既存の Python コードをそのまま利用できます。（便利！！）

違いとしては fn では型定義が必須になります。

# Python
def sum(n, m):
	return n + m

# Mojo
fn sum(n: Int, m: Int) -> Int:
	return n + m

borrowed と inout と owned

Python には無い（ユーザーレベルで定義することは無い）概念として、「borrowed」と「inout」と「owned」という関数の引数渡し制御を明記する必要があります。

• borrowed：不変の引数（関数内で仮引数の変更不可）
  • 引数のデフォルト値になっており、これだけは省略可能

# Mojo
fn sum(borrowed n: Int, borrowed m: Int) -> Int:
	n += 1  # エラーになる
	return n + m

fn main():
  var a = 1
  var b = 2
  let total = sum(a, b)
  print(a)
  print(b)
  print(total)

• inout：可変の引数（参照渡しのため、関数外にも影響）
  • Python のミュータブルオブジェクトと同じ（リストや辞書など）

# Mojo
fn sum(inout n: Int, inout m: Int) -> Int:
	n += 1
	return n + m

fn main():
  var a = 1
  var b = 2
  let total = sum(a, b)
  print(a)  # 2
  print(b)  # 2
  print(total)  # 4

• owned：可変の引数（値渡しのため、関数外には影響しない）
  • Python のイミュータブルオブジェクトと同じ（数値や文字列など）

# Mojo
fn sum(owned n: Int, owned m: Int) -> Int:
	n += 1
	return n + m

fn main():
  var a = 1
  var b = 2
  let total = sum(a, b)
  print(a)  # 1
  print(b)  # 2
  print(total)  # 4

C++ などの言語のように、わざわざ明記する必要があるのはちょっと面倒ですね！（笑）

これによって、安全性が増して速度も上がるのなら、徳ではありますね！（Pythonのシンプルさが恋しい！笑）

構造体

Mojoでは、Python にはない概念として新しく構造体が導入されました！C++などの言語にあるあの構造体です！

記法としては Python の class dataclass と同じ感じですが、動的か静的かの違いがあるため、より実行速度が上がります！

# python
from dataclasses import dataclass

@dataclass
class Person:
    name: str
    age: int

    def __init__(self, name, age):
        self.name = name
        self.age = age

    def show(self):
        print(f'{self.name}は{self.age}歳です')

person1 = Person('太郎', 20)
person1.show()

# Mojo
struct Person:
    var name: String
    var age: Int

    fn __init__(inout self, name: String, age: Int):
        self.name = name
        self.age = age

    fn show(self):
        print(self.name + "は" + self.age + "歳です")

fn main():
	let person1 = Person("太郎", 20)
	person1.show()

注意としては、コンストラクタの１つ目の引数 self には inout として渡す必要があります。なぜなら、inoutは可変な参照を意味し、コンストラクタ内でselfのプロパティを変更できるようにする必要があるからです。また、コンストラクタの外でもselfの変更が反映されるように inout にする必要があります！

Pythonモジュールの利用

Mojo は Python と高い互換性があるため、Python のモジュールを活用することができます！しかし、Python モジュールに関してはその箇所のみ Python のインタプリタを使用するため、使用箇所においてはそこまで実行速度の上昇は見込めないようです。

また、注意点として、関数に raises を示すことが必須になっています。これが無いとエラーになります。

# Python
import numpy as np

array = np.array([1, 2, 3])
print(array)

# Mojo
from python import Python

fn use_array() raises:
    let np = Python.import_module("numpy")
    
    let array = np.array([1, 2, 3])
    print(array)

fn main() raises:
    use_array()

Mojoファイル内に書かなくても、外部の Python ファイルを読み込むこともできます！(部分的な改修や既存コードの追加などにも使えそう！)

カレントディレクトリに mypython.py を作成し、mojo の方でその Python ファイルをモジュールとしてインポートすることができます！

# mypython.py
import numpy as np

def gen_random_values(size, base):
    random_array = np.random.rand(size, size)
    return random_array + base

# main.mojo
from python import Python

fn main() raises:
    Python.add_to_path(".")
    let mypython = Python.import_module("mypython")

    let values = mypython.gen_random_values(2, 3)
    print(values)

エラー：Mojo/Python interoperability error: Unable to locate a suitable libpython, please set MOJO_PYTHON_LIBRARY

Google Colab などの環境では、libpython を読み込めないというエラーが出ます。下記の手順でパスを探し、実行時に環境変数として渡す必要があります！

• libpython までの場所を確認する

!python3 -c 'import sysconfig; print(sysconfig.get_config_var("LIBDIR"))'

• libpython のファイルを確認する

!ls $(python3 -c 'import sysconfig; print(sysconfig.get_config_var("LIBDIR"))') | grep libpython

• 上記2つの結果を合わせて、実行時に環境変数を追加する（３行目を追加）

# Mojoの実行
!export MODULAR_HOME="/root/.modular" && \
export PATH="/root/.modular/pkg/packages.modular.com_mojo/bin:$PATH" && \
export MOJO_PYTHON_LIBRARY="/usr/lib/x86_64-linux-gnu/libpython3.10.so" && \
mojo main.mojo

機械学習による Python との速度比較

Mojo は実行速度や設計思想から、AI 分野に重きを置いて開発された言語ともいわれています。そのため、今回は実際に機械学習による速度の差がどのくらいあるのか検証しました！

🚫 結論：失敗、、、。 Mojo が機械学習系の Python モジュールを全然対応していなく、何も使えないのが現状です。 Python の深層学習フレームワーク「Pytorch」を紹介として載せておきます！

Python による機械学習

ライブラリ等のインストール

!pip install scikit-learn
!pip install numpy
!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

PyTorchによる深層学習

機械学習では定番の「アヤメの分類」の学習モデルです！
PyTorch を用いるとこんなに簡単に深層学習の実装ができます！（すごすぎる！！）

内容は本題ではないためコードの説明は省きますが、一般的な事しかやっていないので、機械学習をしたことがある人なら分かると思います！

from sklearn import datasets
from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import time

start = time.time()

# モデルの定義
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # x(入力)のユニット数は4 
        self.fc1 = nn.Linear(4, 10)
        # 隠れ層1のユニット数は10
        self.fc2 = nn.Linear(10, 10)
        # 隠れ層2のユニット数は10
        self.fc3 = nn.Linear(10, 3)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.log_softmax(self.fc3(x),dim=1)
        return x

# 前処理
iris = datasets.load_iris()  # データセット
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.25 ,random_state=0)
x = torch.tensor(X_train, dtype=torch.float, device="cuda")  # テンソルに変換
y = torch.tensor(y_train, device="cuda")

# 学習の実行
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # GPU 使用
model = Net().to(device)  # 学習モデル
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()  # 損失関数

sum_loss = 0.0 
epoch = 5000
for i in range(1,epoch):
    # 勾配の初期化
    optimizer.zero_grad()
    # 説明変数xをネットワークにかける
    output = model(x)
    # 損失関数の計算
    loss = criterion(output, y)
    # 勾配の計算 
    loss.backward()
    # パラメタの更新
    optimizer.step()

    sum_loss += loss.item()
    if i % 1000 == 0:
        print("loss : {0}".format(sum_loss/i))

# 推論の実行
def result():
  outputs = model(torch.tensor(X_test, dtype=torch.float, device="cuda"))
  _, predicted = torch.max(outputs.data, 1)
  y_predicted = predicted.cpu().numpy() # GPUからCPUに移動
  accuracy = 100 * np.sum(y_predicted == y_test) / len(y_predicted)
  return accuracy

result = result()
print('accuracy: {:.1f}%'.format(result))
end = time.time()
print("実行時間：", end - start)

Mojo による機械学習

Python モジュールを利用して、上記のコードを、型宣言や関数、構造体などの基本的な変換を行って Mojo による実装を試みたが、機械学習系の Python モジュールが全く対応しておらず、インポートでエラーになってしまいます。そもそも構造体には継承の概念がないので、インポートできても現状では無理そうですね。AI 向けの言語みたいな売り方をしているのに、非常に残念ですね。今後のエコシステムに期待ですね！

おまけ

速度比較でおなじみのライプニッツの公式による円周率で速度比較してみました！

# Python
import time

print( "start" )

limit = 10
for m in range(0, 9):
    begin = time.time()
    n = 1
    l = 0.0
    while n < limit:
        l += (1.0 / n) - (1.0 / (n + 2))
        n += 4

    print( l * 4, time.time() - begin, " sec." )
    limit *= 10
print( "end" )

# Mojo
fn main():
  from time import now
  print( "start" )

  var limit = 10
  for m in range(0, 9):
      let start = now()
      var n = 1
      var l = 0.0
      while n < limit:
          l += (1.0 / n) - (1.0 / (n + 2))
          n += 4
      let end = now()
      print( l * 4, (end - start) * 1e-9, " sec." )
      limit *= 10
  print( "end" )

結果はなんと 255 倍！確かに圧倒的に早い！というかあり得ないほど早すぎる！！（笑）

これはすごい言語が出てきましたね！！

	Python	Mojo
時間(9回目)	93.3013961315155	0.36559400600000003
倍率	1	255.20493935974295

終わりに

今回、Python の 68,000 倍も実行速度が速い新言語「Mojo」について、触ってみました！

記法もおおかた Python と近く、Mojo 独自の機能によって安全性や速度が増すのはとても魅力的だと感じました！ただ、正直まだまだ対応しているライブラリや機能が少なく、数年前の初期の「Flutter」よりも全然完成度は低いと私は感じました。まだまだ、使えたような言語ではないですが、今後のアップデートには期待していきたいですね！

ということで、Mojo の紹介でした～！