Mojo言語　〜Tensor型〜

2024年4月7日 21:41

MojoのTensor型にいじめられた

公式のリファレンスを見ながらTensor型で遊んだらよくわからんかった。
と言うことで、今回遊んでみてわかったTensor型の使い方をまとめます。
公式のTensor型のリファレンスは以下で確認できます。
具体例で示したコードはコピペで動くはずです。

ちなみに、下記ページでMojoをすぐに動かして遊べます。

とりあえずTensor型を使う

一旦はTensor型を使ってみよう。

import

まずはtensorをインポートします。

from tensor import Tensor

これでTensor型の値を扱えるようになります。

宣言

以下のコードでTensor型を宣言できます。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    print(x)

結果は以下のように返ってくるはず。

Tensor([[0.0, 0.0, 0.0]], dtype=float32, shape=3)

宣言箇所を詳しく見ると、

var x = Tensor[配列の中身の型](配列の形状)

と言う構成になってます。

要素を取得

これは基本的にPythonと同じようにアクセスできます。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    print(x[0])
    print(x[1])
    print(x[2])

結果は以下の通り。

0.0
0.0
0.0

Tensorのインデックスを指定すれば要素の値を取得できます。

値の代入

代入もPythonと同じ流れで可能です。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    x[0] = 1
    x[1] = 2
    x[2] = 3

    print(x[0])
    print(x[1])
    print(x[2])

結果は以下の通り。

1.0
2.0
3.0

代入できてますね。

Tensor型を便利にするなんやかんや

配列とかを扱う上でやりたくなる操作をTensor型で実行してみます。

num_elements: 要素数の取得

Tensor型の要素数を取得するには、「num_elements」を使います。

num_elements(self: Self) -> Int

戻り値はInt型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    print(x.num_elements())

これで要素数を取得できます。

おまけ: 要素を順番に取り出す

Pythonだとlist型等はそのままでも順番に取り出せますが、
MojoのTensor型は要素を順番に取り出す機能を持っていません。
なので、for文でインデックスを指定して順番に要素を取り出します。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    var n = x.num_elements()
    for i in range(n):
        print(x[i])

type: 型の確認

宣言したTensor型が保持している型を確認するときは「type」を使います。

type(self: Self) -> DType

戻り値はDType型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    print(x.type())

結果は以下の通り。

float32

astype: Tensorの要素の型変換

ここまでfloat32でTensor型を使ってますが、他の型に変換する時は
「astype」を使います。

astype[new_dtype: DType](self: Self) -> Tensor[new_dtype]

戻り値はTensor型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](3)
    var y = x.astype[DType.uint8]()
    print(y)

結果は以下の通り。

 Tensor([[0, 0, 0]], dtype=uint8, shape=3)

型変換できてますね。

clip: Tensorの値を上下限で切る

上下限を設定し、Tensor型の要素の値の中から上下限からはみ出した値は
自動的に上限または下限の値に置換される機能が「clip」です。

clip(self: Self, lower_bound: SIMD[dtype, 1], upper_bound: SIMD[dtype, 1]) -> Self

戻り値はTensor型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](5)
    var lower = 3
    var upper = 6
    for i in range(x.num_elements()):
        x[i] = i ** 2

    var c = x.clip(lower, upper)
    print(x)
    print(c)

結果は以下の通り。

Tensor([[0.0, 1.0, 4.0, 9.0, 16.0]], dtype=float32, shape=5)
Tensor([[3.0, 3.0, 4.0, 6.0, 6.0]], dtype=float32, shape=5)

Tensorにはi^2の値を格納しています。
下限を3.0、上限は6.0に設定してxにclipを実行すると、
下限未満の値は下限値に、上限を超える値は上限値に置換されます。

data: ポインターを取得

Tensorが格納されているポインターを取得する時は「data」を使います。

data(self: Self) -> DTypePointer[dtype, 0]

戻り値はDTypePointer型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](5)
    print(x.data())

これでポインターが返ってきます。

rank: Tensorの階数を取得

Tensorの階数を取得する時は「rank」を使います。

rank(self: Self) -> Int

戻り値はInt型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](5)
    print(x.rank())

bytecount: Tensorのサイズを取得

Tensor型のデータサイズを取得するときは「bytecount」を使います。

bytecount(self: Self) -> Int

戻り値はInt型。
実際に使うコードは以下の通り。

from tensor import Tensor


fn main():
    var x = Tensor[DType.float32](5)
    var y = Tensor[DType.float32](5)
    var z = Tensor[DType.float32](100)
    var x_i = Tensor[DType.uint8](5)
    var y_i = Tensor[DType.uint8](100)
    print(x.bytecount())
    print(y.bytecount())
    print(z.bytecount())
    print(x_i.bytecount())
    print(y_i.bytecount())

結果は以下の通り。

サイズも型も同じxとyはサイズが同じになっています。
float32のTensorは初期状態で1要素あたり4Byte、
uint8のTensorは初期状態で1要素あたり1Byte使ってるみたいですね。

Tensorをもっと便利に扱いたい！

まだまだあります、Tensorの機能。
さらにインポートを増やしてできることも増やしましょう。

TensorSpec型

Tensor型の構造を保持する型です。
以下のインポートを追加します。

from tensor import TensorSpec

TensorSpecを使うと、Tensor型の宣言と具体的なTensorの構造を
分離して扱えるようになって便利ですね。

具体的には以下のように使用します。

from tensor import Tensor, TensorSpec


fn main():
    var width = 5
    var height = 5
    var dtype = DType.float32
    
    var spec = TensorSpec(dtype, height, width)
    var x = Tensor[DType.float32](spec)
    
    print(x)

結果は以下の通り。


Tensor([[0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0]], dtype=float32, shape=5x5)

5x5のTensorができました。
以下のようにすることで、カラー画像形式のTensorも作れます。

from tensor import Tensor, TensorSpec


fn main():
    var width = 5
    var height = 5
    var channel = 3
    var dtype = DType.uint8
    
    var spec = TensorSpec(dtype, height, width, channel)
    var x = Tensor[DType.uint8](spec)
    
    print(x)

TensorSpecに「高さx幅xチャンネル」を指定しました。
結果は以下の通り。

Tensor([[[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0]],
[[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0]],
[[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0]],
[[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0]],
[[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0],　[0, 0, 0]]],
dtype=uint8, shape=5x5x3)

見やすくするために少し整形しましたが、
5x5の各要素に3つの値を保持した配列がさらに入っていますね。

TensorShape型

Tensorの形状を保持する型です。
TensorShape型を使うには、以下のインポートを追加します。

from tensor import TensorShape

Tensorのリサイズの時に使ったりします。
TensorShapeの具体的な宣言方法は以下の通り。

from tensor import Tensor, TensorShape


fn main():
    var width = 5
    var height = 5
    var channel = 3
    
    var shape = TensorShape(height, width, channel)
    print(shape)

結果は以下の通り。

5x5x3

形状がTensorShape型で出力されました。

TensorSpecとTensorShapeで使える機能

ここまで解説した以下の機能はTensorSpecとTensorShapeでも使えます。

rank
num_elements

TensorSpecのみ(TensorShapeでは使えない)機能には以下があります。

dtype
bytecount

実際に使ってみましょう。

from tensor import Tensor, TensorSpec, TensorShape


fn main():
    var width = 5
    var height = 5
    var channel = 3
    var dtype = DType.uint8
    var spec = TensorSpec(dtype, height, width, channel)
    var shape = TensorShape(height, width, channel)
    
    print('TensorSpec.rank = ' + str(spec.rank()))
    print('TensorSpec.num_elements = ' + str(spec.num_elements()))
    print('TensorSpec.dtype = ' + str(spec.dtype()))
    print('TensorSpec.bytecount = ' + str(spec.bytecount()) + '\n')
    
    print('TensorShape.rank = ' + str(shape.rank()))
    print('TensorShape.num_elements = ' + str(shape.num_elements()))

結果は以下の通り。

TensorSpec.rank = 3
TensorSpec.num_elements = 75
TensorSpec.dtype = uint8
TensorSpec.bytecount = 75

TensorShape.rank = 3
TensorShape.num_elements = 75

それぞれの具体的な動き方については変わりませんね。

Tensor + TensorSpec + TensorShape

さらに便利な機能がまだあるので、それらも見ていきましょう。

ここから使うTensorのテンプレート

以降はTensorの値を初期値ではなくちゃんと値を入れて扱います。
多分その方がわかりやすいので。
ただ、毎回ごちゃごちゃ中身を変えるとややこしいので、
決まった中身のTensorを先に決めておきます。

from tensor import Tensor, TensorSpec


fn main():
    var count = 225
    var dtype = DType.uint8
    var spec = TensorSpec(dtype, count)
    var template = Tensor[DType.uint8](spec)
    for i in range(template.num_elements()):
        template[i] = 255 - 17 * i

このTensorを使いますが、main関数が見にくくなるので外に出します。

from tensor import Tensor, TensorSpec


fn templates() -> Tensor[DType.uint8]:
    var count = 225
    var dtype = DType.uint8
    var spec = TensorSpec(dtype, count)
    var template = Tensor[DType.uint8](spec)
    for i in range(template.num_elements()):
        template[i] = 255 - 17 * i
    return template


fn main():
    var template = templates()
　　　　　　　　print(template)

templates関数の戻り値が先ほど決めたTensorになっています。

結果は以下の通り。

Tensor([[255, 238, 221, ..., 65, 48, 31]], dtype=uint8, shape=225)

このTensorをベースとして、結果の変化に注目しましょう。

以降はこのtemplate関数を省略して書きますのでご注意ください。

ireshape: Tensorの形状変更

「ireshape」を使うとTensorの形状を任意の形に変形することができます。
ireshapeの形状はTensorShape型を使用して指定します。

from tensor import Tensor, TensorSpec, TensorShape


fn main():
    var template = templates()
    var shape = TensorShape(15, 15)
    try:
        template.ireshape(shape)
    except:
        print('RESHAPE ERROR !!')
    print(template)

結果は以下の通り。

Tensor([[255, 238, 221, ..., 51, 34, 17],
[0, 239, 222, ..., 52, 35, 18],
[1, 240, 223, ..., 53, 36, 19],
..., 
[11, 250, 233, ..., 63, 46, 29],
[12, 251, 234, ..., 64, 47, 30],
[13, 252, 235, ..., 65, 48, 31]], dtype=uint8, shape=15x15)

Tensorが15x15に形状が変わってますね。
注意点として、指定した形状と元のTensorの要素数で整合性が取れなかった
場合にMojoはエラーを送出するので、tryで囲む必要があります。

reshape: 形状変更した新しいTensorを作成する

先ほどのireshapeは元のTensorを形状変更するため、
元のTensorを保持しない破壊的な操作ですとなります。
元のTensorを保持した上で新しい形状のTensorが欲しい場合は、
「reshape」を使います。

from tensor import Tensor, TensorSpec, TensorShape


fn main():
    var template = templates()
    var shape = TensorShape(15, 15)
    try:
        var c = template.reshape(shape)
        print(c)
    except:
        print('RESHAPE ERROR !!')
    print(template)

結果は以下の通り。

Tensor([[255, 238, 221, ..., 51, 34, 17],
[0, 239, 222, ..., 52, 35, 18],
[1, 240, 223, ..., 53, 36, 19],
..., 
[11, 250, 233, ..., 63, 46, 29],
[12, 251, 234, ..., 64, 47, 30],
[13, 252, 235, ..., 65, 48, 31]], dtype=uint8, shape=15x15)
Tensor([[255, 238, 221, ..., 65, 48, 31]], dtype=uint8, shape=225)

元のTensorが保持された状態で新しい形状のTensorを生成できました。

spec: Tensorのspecを取得

TensorのspecをTensorSpec型で取得するのが「spec」です。

from tensor import Tensor, TensorSpec, TensorShape


fn main():
    var template = templates()
    var shape = TensorShape(15, 5, 3)
    
    try:
        template.ireshape(shape)
    except:
        print('RESHAPE ERROR !!')
        
    print(template.spec())

結果は以下の通り。

15x5x3xuint8

Tensorのspecが取得できました。
戻り値の型はTensorSpec型です。

shape: Tensorのshapeを取得

Tensorの形状をTensorShape型で取得するのが「shape」です。

from tensor import Tensor, TensorSpec, TensorShape


fn main():
    var template = templates()
    var shape = TensorShape(15, 5, 3)
    
    try:
        template.ireshape(shape)
    except:
        print('RESHAPE ERROR !!')
        
    print(template.shape())

結果は以下の通り。

15x5x3

Tensorの形状が取得できました。
戻り値はTensorShape型です。

dim: Tensorの次元を取得

Tensor内の特定のインデックスの次元を取得するには「dim」を使います。

from tensor import Tensor, TensorSpec, TensorShape


fn main():
    var template = templates()
    var shape = TensorShape(15, 5, 3)
    
    try:
        template.ireshape(shape)
    except:
        print('RESHAPE ERROR !!')
        
    print(template.dim(0))
    print(template.dim(1))
    print(template.dim(2))

結果は以下の通り。

15
5
3

TensorShapeのそれぞれの値を取得するには便利そうですね。
戻り値はInt型です。

おわりに

それなりに使う機能はひとまずまとめることができたかな〜と思います。
基本的に備忘録的な立ち位置なので、不備とかわけわからんことも書いてる可能性があります。
ご容赦願います。

この記事が気に入ったらサポートをしてみませんか？

Mojo言語 〜Tensor型〜

MojoのTensor型にいじめられた

とりあえずTensor型を使う

import

宣言

要素を取得

値の代入

Tensor型を便利にするなんやかんや

num_elements: 要素数の取得

おまけ: 要素を順番に取り出す

type: 型の確認

astype: Tensorの要素の型変換

clip: Tensorの値を上下限で切る

data: ポインターを取得

rank: Tensorの階数を取得

bytecount: Tensorのサイズを取得

Tensorをもっと便利に扱いたい！

TensorSpec型

TensorShape型

TensorSpecとTensorShapeで使える機能

Tensor + TensorSpec + TensorShape

ここから使うTensorのテンプレート

ireshape: Tensorの形状変更

reshape: 形状変更した新しいTensorを作成する

spec: Tensorのspecを取得

shape: Tensorのshapeを取得

dim: Tensorの次元を取得

おわりに

Mojo言語　〜Tensor型〜