見出し画像

ゼロから始める自作LLM

イントロダクション

今回は、オズの魔法使いの文章をベースに自作LLMをGoogle Colabで作成していきます。LLMのベースとなる学習データがオズの魔法使いという非常に少ない文章量のため、実行結果は期待されているほどにはなりません。しかし、ゼロからLLMを作成してみるという経験を積むには良いと考えてます。

今回は、こちらのコードを参考にさせて頂きました。


コード説明

今回は、Google ColabのA100かV100で実行します。それ以外だと計算に時間がかかりすぎてしまいます。

実行する前に、上記のGitHubでwizard_of_oz.txtをダウンロードして/contentに配置しましょう。

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

PyTorchライブラリをインストールします。

device = 'cuda' if torch.cuda.is_available() else 'cpu'

ここでは、CUDA(NVIDIAのGPUを使用するためのプラットフォーム)が利用可能かどうかを確認しています。CUDAが利用可能な場合はcudaをデバイスとして選択し、利用できない場合はcpuを使用します。これにより、計算速度の向上を図ることができます。

wizard_of_oz_book = "wizard_of_oz.txt"

with open(wizard_of_oz_book, 'r', encoding='utf-8') as f:
    text = f.read()
    chars = sorted(list(set(text)))

vocab_size = len(chars)

テキストファイルを読み込んで、そのテキスト内の異なる文字の種類をカウントし、語彙サイズを計算する処理を行っています。

string_to_int = { ch:i for i, ch in enumerate(chars) }
int_to_string = { i:ch for i, ch in enumerate(chars) }
encode = lambda s: [string_to_int[c] for c in s]
decode = lambda l: ''.join([int_to_string[i] for i in l])

data = torch.tensor(encode(text), dtype=torch.long)
data[:100] # First 100 characters of book tokonized

テキストデータを整数に変換し、その逆の操作も行うマッピングを作成し、PyTorchのテンソルに変換する処理を実装しています。

マッピングの作成:

string_to_int = { ch:i for i, ch in enumerate(chars) }: この辞書内包表記は、各文字を対応するインデックスにマッピングします。enumerate(chars)は各文字にインデックスを付与し、それを辞書として格納しています。
int_to_string = { i:ch for i, ch in enumerate(chars) }: こちらはその逆の操作で、インデックスから文字を取得できるようにマッピングしています。


エンコードとデコードの関数:

encode = lambda s: [string_to_int[c] for c in s]: 与えられた文字列sを、string_to_intマッピングを使用して整数のリストに変換するラムダ関数です。
decode = lambda l: ''.join([int_to_string[i] for i in l]): 整数のリストlを文字列に戻すためのラムダ関数です。int_to_stringマッピングを使用しています。

PyTorchのテンソルへの変換:

data = torch.tensor(encode(text), dtype=torch.long): ここでencode関数を使用してテキストをエンコードし、その結果をPyTorchのテンソルに変換します。dtype=torch.longは、生成されるテンソルのデータタイプが長整数型であることを指定しています。

テンソルのサブセットの表示:

data[:100]: ここではテンソルdataの最初の100要素を取得します。これは、テキストの最初の100文字がどのように整数に変換されたかを示すものです。

# Let's now split up the data into train and validation sets
n = int(0.9*len(data))
train_data = data[:n]
val_data = data[n:]

テキストデータをエンコードした後のPyTorchのテンソルを訓練データセットと検証データセットに分割する操作を行っています。

  1. データセットのサイズの計算:

    • n = int(0.9*len(data)): この式では、dataテンソルの長さ(len(data))の90%に相当する位置を整数で計算しています。このnは、訓練データセットのデータ点の数を示します。

  2. 訓練データと検証データの分割:

    • train_data = data[:n]: データの最初からn番目の要素までを訓練データセットとして切り出しています。これにより、データの90%が訓練用に確保されます。

    • val_data = data[n:]: n番目の要素からデータの末尾までを検証データセットとして切り出しています。残りの10%が検証用に使用されます。

batch_size = 32
block_size = 128

def get_batch(split):
    # generate a small batch of data of inputs x and targets y
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

xb, yb = get_batch('train')
print('inputs:')
print(xb.shape)
print(xb)
print('targets:')
print(yb.shape)
print(yb)

与えられた訓練データセットまたは検証データセットから、小さなバッチの入力データ x とターゲットデータ y を生成するための関数 get_batch を定義し、実際に訓練データセットからバッチを取得して内容を表示しています。

関数 get_batch の定義:

  1. パラメータ:

    • split: 'train' または 'val' のいずれかの文字列を受け取り、訓練データか検証データかを選択します。

  2. データの選択:

    • data = train_data if split == 'train' else val_data: split の値に応じて、train_data または val_data を選択します。

  3. ランダムインデックスの生成:

    • ix = torch.randint(len(data) - block_size, (batch_size,)): データセットからランダムに開始位置を選ぶためのインデックスを生成します。block_size をデータの長さから引くのは、選択したインデックスから block_size 分のデータを安全に取得できるようにするためです。

  4. バッチデータの生成:

    • x = torch.stack([data[i:i+block_size] for i in ix]): 各インデックス i から block_size の長さのデータを取得してスタックします。これがモデルの入力データ x になります。

    • y = torch.stack([data[i+1:i+block_size+1] for i in ix]): 各インデックスの次の位置から block_size の長さのデータを取得してスタックします。これが正解ラベル(ターゲットデータ y)になります。

  5. デバイスへのデータの転送:

    • x, y = x.to(device), y.to(device): 入力 x とターゲット y を事前に設定されたデバイス(GPUまたはCPU)に転送します。

バッチデータの取得と表示:

  • xb, yb = get_batch('train'): get_batch 関数を呼び出し、訓練データからバッチを取得します。

  • print ステートメントを使用して、取得したバッチの形状と内容を表示します。

コードの出力:

  • print('inputs:'): 入力データ xb の形状と内容を表示します。

  • print('targets:'): ターゲットデータ yb の形状と内容を表示します。

n_embd = 384
n_head = 8
n_layer = 8
dropout = 0.2

ディープラーニングモデル(特にトランスフォーマー系のモデル)のパラメータを定義しています。

  1. n_embd = 384:

    • これはモデルの各レイヤーで使用される埋め込み(embedding)のサイズです。埋め込みサイズは、モデルが扱う特徴の次元数を表し、ここでは384次元と設定されています。大きい埋め込みサイズは、より豊富な情報をキャプチャする能力が高まりますが、計算コストも増大します。

  2. n_head = 8:

    • トランスフォーマーモデルのマルチヘッド・アテンション機構で使用される「ヘッド」の数です。ここでは8ヘッドが使用されています。各ヘッドは異なる視点から情報を処理し、モデルが異なるサブスペースから特徴を抽出するのを助けます。

  3. n_layer = 8:

    • これはモデルの層の数、つまり何層のエンコーダーまたはデコーダー層を使用するかを指定します。ここでは8層が使用されています。層の数が多いほどモデルは複雑な特徴や関係を学習する能力が高まりますが、同時に過学習のリスクや訓練時間の増加も考慮する必要があります。

  4. dropout = 0.2:

    • ドロップアウト率で、モデル訓練時にランダムに選ばれたノード(ニューロン)を無効化(0に設定)する確率を示します。ここでは20%の確率でニューロンがドロップアウトされます。ドロップアウトは一般的にモデルの過学習を防ぐために使用され、より堅牢なネットワーク学習を促進します。

class Head(nn.Module):
    # one head of self-attention
    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B,T,C = x.shape
        k = self.key(x)   # (B,T,C)
        q = self.query(x) # (B,T,C)
        # compute attention scores ("affinities")
        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)
        wei = F.softmax(wei, dim=-1) # (B, T, T)
        wei = self.dropout(wei)
        # perform the weighted aggregation of the values
        v = self.value(x) # (B,T,C)
        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)
        return out

トランスフォーマーモデルのための自己注意メカニズムの一部分を実装するためのHeadクラスを定義しています。このクラスはnn.Moduleを継承しており、自己注意の1つのヘッドを表現します。以下、クラスとそのメソッドの詳細な説明をします。

クラス Head の定義:

  1. __init__メソッド:

    • head_size: このヘッドが処理する特徴のサイズを指定する引数です。

    • self.key, self.query, self.value: それぞれ自己注意メカニズムにおいてキー、クエリ、バリューの計算に使われる線形層です。これらは同じ入力次元n_embdを持ち、出力次元はhead_sizeです。バイアスは使用されていません。

    • self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size))): 下三角行列をバッファとして登録しており、マスキング処理に使用します。これは注意スコアの計算で将来の情報が使われないようにするためです。

    • self.dropout: ドロップアウト層で、初期化時に設定したdropout値を使用しています。

  2. forwardメソッド:

    • 入力x(B, T, C)形式のテンソルです。ここでBはバッチサイズ、Tは系列長、Cは特徴量の次元数です。

    • k, q, v: 入力xをそれぞれキー、クエリ、バリュー線形層に通すことで、対応する変換が行われます。

    • wei = q @ k.transpose(-2, -1) * C**-0.5: クエリとキーのトランスポーズとの行列積を計算し、スケーリングを行います。

    • マスキング処理: wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))で、不要な将来の情報を含む部分を無限大の負の値で置き換えます(ソフトマックス前に)。

    • ソフトマックス関数を適用して、正規化された注意スコアweiを得ます。

    • ドロップアウトを適用した後、注意スコアを使用して値ベクトルvの加重和を計算し、出力outを得ます。

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

ランスフォーマーモデル内で重要な役割を果たすマルチヘッド注意(MultiHead Attention)のメカニズムを実装するためのMultiHeadAttentionクラスを定義しています。このクラスもnn.Moduleを継承しており、複数の自己注意のヘッド(Headクラスのインスタンス)を組み合わせて使用します。

クラス MultiHeadAttention の定義:

  1. __init__メソッド:

    • num_heads: 使用するヘッドの数を指定します。

    • head_size: 各ヘッドが処理する特徴のサイズを指定します。

    • self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)]): 指定された数のヘッドを作成し、ModuleListに格納します。これにより、各ヘッドが個別に異なるサブスペースで情報を処理することができます。

    • self.proj = nn.Linear(n_embd, n_embd): 全てのヘッドの出力を結合した後に、この線形層を通じて元の埋め込みサイズに射影します。これは、ヘッドの出力を適切に統合するための重要なステップです。

    • self.dropout = nn.Dropout(dropout): ドロップアウト層で、初期化時に設定したdropout値を使用して、過学習を防ぎます。

  2. forwardメソッド:

    • 入力xに対して、各ヘッドの自己注意処理を行い、その結果をリスト内包表記で取得します。torch.cat([h(x) for h in self.heads], dim=-1)は、各ヘッドの出力を最後の次元(特徴の次元)に沿って結合します。これにより、異なる視点からの情報が一つの大きなテンソルに統合されます。

    • 結合された出力はself.proj(out)を通して再び元の特徴サイズに射影され、その後self.dropoutでドロップアウトを適用します。

class FeedFoward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

トランスフォーマーモデル内で利用される位置ごとの前方フィードフォワードネットワーク(Feed Forward Network、FFN)を実装するためのFeedForwardクラスを定義しています。nn.Moduleを継承し、ニューラルネットワークのレイヤーをnn.Sequentialを使ってシーケンシャルに組み立てています。

クラス FeedForward の定義:

  1. __init__メソッド:

    • n_embd: 入力および出力の特徴次元数を指定します。この値はトランスフォーマーの埋め込みサイズと一致します。

    • self.net: このシーケンシャルコンテナ内に以下の層が組み込まれています。

      • nn.Linear(n_embd, 4 * n_embd): 最初の線形層で、入力次元をn_embdから4 * n_embdに拡大します。これにより、ネットワークはより高次の特徴を抽出する能力を持ちます。

      • nn.ReLU(): 非線形活性化関数で、モデルの表現力を向上させます。

      • nn.Linear(4 * n_embd, n_embd): 拡大された特徴空間から元の埋め込みサイズに戻すための線形層です。

      • nn.Dropout(dropout): ドロップアウト層で、過学習を防ぐために設定されたドロップアウト率を用います。

  2. forwardメソッド:

    • 入力xを受け取り、self.netを通して一連の変換を施した後、結果を返します。x(batch_size, sequence_length, n_embd)の形状を持つテンソルと想定されます。

class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        # n_embd: embedding dimension, n_head: the number of heads we'd like
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedFoward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

トランスフォーマーモデルの一部である基本的なブロック構造(Blockクラス)を定義しています。このクラスはnn.Moduleを継承し、マルチヘッド自己注意機構と前方フィードフォワードネットワークを含んでいます。

クラス Block の定義:

  1. __init__メソッド:

    • n_embd: 埋め込みの次元数です。

    • n_head: 使用するヘッドの数です。

    • head_size = n_embd // n_head: 各自己注意ヘッドのサイズを計算します。全体の埋め込み次元をヘッドの数で割ることにより、各ヘッドが扱う特徴の次元数を決定します。

    • self.sa = MultiHeadAttention(n_head, head_size): マルチヘッド自己注意層を初期化します。

    • self.ffwd = FeedFoward(n_embd): 前方フィードフォワードネットワーク層を初期化します。

    • self.ln1 = nn.LayerNorm(n_embd): 自己注意層の前に適用されるレイヤー正規化です。

    • self.ln2 = nn.LayerNorm(n_embd): 前方フィードフォワード層の前に適用されるレイヤー正規化です。

  2. forwardメソッド:

    • 入力xに対して、次の操作を順番に適用します:

      • self.ln1(x): 最初にレイヤー正規化を行います。

      • self.sa(self.ln1(x)): 正規化した出力をマルチヘッド自己注意層に通します。

      • x = x + self.sa(self.ln1(x)): 入力と自己注意の出力を足し合わせます(残差接続)。

      • self.ln2(x): 次に再びレイヤー正規化を行います。

      • self.ffwd(self.ln2(x)): 正規化した出力を前方フィードフォワードネットワークに通します。

      • x = x + self.ffwd(self.ln2(x)): 再び入力とフィードフォワードの出力を足し合わせます(残差接続)。

class GPTLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd) # final layer norm
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.apply(self._init_weights)

    # Initialize weights and bias
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, index, targets=None):        
        B, T = index.shape

        tok_emb = self.token_embedding_table(index)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
        x = tok_emb + pos_emb # (B, T, C)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x) # (B, T, vocab_size)
        
        if targets is None:
            loss = None
        else:                        
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, index, max_new_tokens):
        for _ in range(max_new_tokens):
            index_cond = index[:, -block_size:]
            
            # get the predictions
            logits, loss = self.forward(index_cond)
            
            # focus only on the last time step
            logits = logits[:, -1, :] # becomes (B, C)
            
            # apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1) # (B, C)
            
            # sample from the distribution
            index_next = torch.multinomial(probs, num_samples=1) # (B, 1)
            
            # append sampled index to the running sequence
            index = torch.cat((index, index_next), dim=1) # (B, T+1) 
            
        return index

PyTorchを使用してGPTスタイルの言語モデルを定義するGPTLanguageModelクラスを実装しています。このクラスはnn.Moduleを継承しており、主にトークンの埋め込み、位置の埋め込み、複数のトランスフォーマーブロック、最終層正規化、および言語モデルのヘッドから構成されています。さらに、重みの初期化、フォワードパス、テキスト生成のためのメソッドが含まれています。

クラス GPTLanguageModel の定義:

  1. __init__メソッド:

    • vocab_size: 語彙のサイズです。

    • self.token_embedding_table: トークン埋め込みテーブルで、各トークンをn_embd次元のベクトルにマッピングします。

    • self.position_embedding_table: 位置埋め込みテーブルで、各位置をn_embd次元のベクトルにマッピングします。

    • self.blocks: トランスフォーマーモデルの主要な処理ブロックを複数持ち、それぞれが自己注意と前方フィードフォワードネットワークから構成されています。

    • self.ln_f: 最終層の正規化を行います。

    • self.lm_head: 最終的な出力を語彙サイズの次元に変換する線形層です。

    • self.apply(self._init_weights): 各層の重みを初期化するために、定義された重み初期化関数を適用します。

  2. 重みの初期化 (_init_weightsメソッド):

    • 線形層と埋め込み層に対して特定の初期化を行います。線形層の重みは正規分布を使用して初期化し、埋め込み層も同様にします。

  3. forwardメソッド:

    • 入力indexと任意のtargetsを受け取り、トークンと位置の埋め込みを行い、トランスフォーマーブロックを通過させ、最終層正規化と線形層で処理を行います。

    • 目的変数targetsが与えられた場合、クロスエントロピー損失を計算します。

  4. generateメソッド:

    • 与えられた初期インデックス(トークンIDの系列)から新しいトークンを生成するメソッドです。

    • 最新のトークンを考慮して次のトークンを予測し、サンプリングを行います。この過程は指定されたトークン数だけ繰り返されます。

m = GPTLanguageModel(vocab_size)
model = m.to(device)

定義されたGPTLanguageModelクラスのインスタンスを作成し、そのモデルを特定の計算デバイス(GPUまたはCPU)に配置しています。

モデルのインスタンス化:

  • m = GPTLanguageModel(vocab_size): GPTLanguageModelクラスから新しいオブジェクトmを作成します。ここでvocab_sizeは、モデルが使用する語彙のサイズを指定する引数です。このサイズはモデルのトークン埋め込み層と言語モデルヘッドで使用されます。

モデルをデバイスに配置:

  • model = m.to(device): .to(device)メソッドを使用して、モデルmを変数deviceで指定されたデバイス(通常はGPUまたはCPU)に配置します。この操作により、モデルのすべてのパラメータ(重みやバイアスなど)が指定されたデバイスに移動され、そのデバイスでの計算が可能になります。

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
steps = 3000

GPTLanguageModelモデル用のオプティマイザを設定し、トレーニングステップの数を指定しています。

オプティマイザの設定:

  • optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4): AdamWオプティマイザを使用してモデルのパラメータを更新するための設定を行っています。ここでmodel.parameters()GPTLanguageModelのすべての学習可能なパラメータを取得します。lr=3e-4は学習率を指定しており、これはオプティマイザが各トレーニングステップでパラメータをどの程度更新するかを決定します。

トレーニングステップの数の設定:

  • steps = 3000: このモデルをトレーニングする際のステップ(イテレーション)の総数を3000と設定しています。これは、オプティマイザがモデルのパラメータを更新する回数を意味します。

AdamWオプティマイザ:

  • AdamWAdamオプティマイザの一種で、ウェイトの減衰(L2正則化の一種)を改善しており、特に深層学習モデルの一般化性能を向上させるために設計されています。これは特に大規模なモデルや過学習が懸念される場合に有効です。

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(250)
        for k in range(250):
            X, Y = get_batch(split)
            logits, loss = model.forward(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

訓練(train)データセットと検証(validation)データセットについて、モデルの損失を推定するためのものです。

  1. @torch.no_grad(): このデコレーターは、PyTorchの計算グラフにおいて、勾配計算を行わないことを指定します。これにより、メモリ使用量が減り、計算速度が向上します。モデル評価時にはパラメータを更新しないため、通常このデコレーターを使用します。

  2. def estimate_loss(): estimate_lossという名前の関数を定義しています。

  3. out = {}: 空の辞書を作成しています。これは、各データセット('train'と'val')の平均損失を格納するために使用されます。

  4. model.eval(): モデルを評価モードに設定します。これは、訓練モードとは異なり、ドロップアウトやバッチ正規化などの層が評価モードで動作するようになります。

  5. for split in ['train', 'val']: 'train'と'val'の各データセットに対してループを行います。

  6. losses = torch.zeros(250): 250要素のゼロで初期化されたテンソルを作成します。これは各バッチの損失を格納するために使用します。

  7. for k in range(250): 各データセットについて250回のバッチ処理を行うループです。

  8. X, Y = get_batch(split): get_batch関数を呼び出し、現在のsplit('train'または'val')に対応するバッチデータを取得します。

  9. logits, loss = model.forward(X, Y): モデルに入力データXと正解データYを渡して前方計算を行い、出力(logits)と損失(loss)を取得します。

  10. losses[k] = loss.item(): 計算された損失のスカラー値を取得し、lossesテンソルに格納します。

  11. out[split] = losses.mean(): 各データセットの250バッチの損失の平均値を計算し、out辞書に格納します。

  12. model.train(): モデルを再び訓練モードに戻します。

  13. return out: 各データセットの平均損失が格納された辞書を返します。

for step in range(steps):
    if step % 250 == 0:
        losses = estimate_loss()
        print(f"step: {step}, train loss: {losses['train']:.3f}, val loss: {losses['val']:.3f}")
    
    # sample a batch of data
    xb, yb = get_batch('train')

    # evaluate the loss
    logits, loss = model.forward(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

print(loss.item())

モデルの訓練プロセスを実行する際の主なステップを示しています。

  1. for step in range(steps)::訓練を行うステップ数をstepsで指定し、その範囲でループを行います。

  2. if step % 250 == 0::ステップ数が250の倍数の時(すなわち、250ステップごとに)、以下の処理を実行します。

    • losses = estimate_loss()estimate_loss関数を呼び出して、現在のモデルの訓練と検証データセットにおける損失を計算します。

    • print(f"step: {step}, train loss: {losses['train']:.3f}, val loss: {losses['val']:.3f}"):計算された損失を表示します。ここで.3fは、小数点以下3桁までのフォーマットを指定しています。

  3. xb, yb = get_batch('train')get_batch関数を使用して訓練データセットからバッチデータを取得します。

  4. logits, loss = model.forward(xb, yb):モデルにバッチデータを入力して前方計算を行い、出力(logits)と損失(loss)を取得します。

  5. optimizer.zero_grad(set_to_none=True):オプティマイザーの勾配をゼロにリセットします。set_to_none=Trueを指定することで、勾配のメモリをNoneに設定し、少しのパフォーマンス向上が期待できます。

  6. loss.backward():損失関数に基づいて、モデルのパラメータに対する勾配を計算します。

  7. optimizer.step():計算された勾配を使用して、モデルのパラメータを更新します。

  8. print(loss.item()):最終ステップの損失値を表示します。

prompt = "I am "
context = torch.tensor(encode(prompt), dtype=torch.long, device=device)
generated_chars = decode(model.generate(context.unsqueeze(0), max_new_tokens=500)[0].tolist())
print(generated_chars)

特定のプロンプトを使用してテキスト生成モデルを起動し、生成されたテキストを表示するプロセスを示しています。

  1. prompt = "I am":生成したいテキストの初めの部分(プロンプト)を指定しています。

  2. context = torch.tensor(encode(prompt), dtype=torch.long, device=device):プロンプトをモデルが理解できる形式にエンコードし、その結果をPyTorchのテンソルに変換しています。ここで、dtype=torch.longはテンソルのデータ型を指定しており、device=deviceは計算を行うデバイス(CPUまたはGPU)を指定しています。

  3. generated_chars = decode(model.generate(context.unsqueeze(0), max_new_tokens=500)[0].tolist()):エンコードされたテンソルをモデルに渡してテキストを生成します。context.unsqueeze(0)はバッチ次元を追加しています(多くのモデルはバッチ処理を前提としているため)。max_new_tokens=500は生成するトークンの最大数を指定しています。生成された結果は、decode関数を使用して人間が読めるテキストにデコードされます。

  4. print(generated_chars):生成されたテキストを表示します。


実行結果

I am mysert, wetched out of a gave betch I'm a humbug Wizard as you drivended in your wings, because I did here was not spires; for my _milk; you must see you ate that all you are lest diamos you are Amiceau." The Wizard explose a servantdoris land enormound about forme in these of yests of much hanging mad charm hairs above them so that even the served sent to life a girl set who had to them live to them tie yet," called the Wizard, "or the Prince look in this long and face some distinguis me to r

実行結果


自作LLM作成に興味がございましたら、他にも自作LLMの記事がございますので見て頂ければ幸いです。


この記事が参加している募集

AIとやってみた

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