Пишем свой PyTorch на NumPy. ФИНАЛ. Запускаем GPT-2
PyTorch — это мощный и гибкий фреймворк для машинного обучения, широко используемый для создания нейронных сетей. Он особенно популярен благодаря простоте использования, динамическим вычислительным графам и богатой экосистеме инструментов для обучения моделей. Для использования этого фреймворка, часто достаточно поверхностно понимать работу алгоритмов машинного обучения.
Но Андрей Карпаты, известный исследователь в области ИИ, считает, что реализация алгоритмов с нуля позволяет понять их суть и детали работы, что сложно осознать, используя только готовые библиотеки. Это помогает развить интуицию для дальнейшего применения и улучшения методов. Андрей посвящает много собственного времени, чтобы объяснять ключевые принципы работы нейросетей в своих блогах и на своём ютуб-канале. Он также не раз подчеркивал, что на его курсе в Cтэнфорде есть задачи по реализации различных алгоритмов, например, обратное распространение.
Я хотел бы посвятить данную статью этой идеи, потому что мне самому особенно интересно копаться в алгоритмах глубокого обучения. Эта статья продолжение второй статьи
Итак, как я и сказал в предыдущей статье, у нас достаточно знаний, чтобы собрать их в целую библиотеку. Так я и сделал, реализовав pycandle.
Там находиться всё то, что мы изучали на протяжении 3 частей. В этой части мы будем писать инференс код для GPT2 на собственной библиотеке!
Начнём с импорта библиотек
import torch
import torch.nn as nn
Ой, не то. Я имел в виду
import candle
import candle.nn as nnТакже воспользуемся токенизатором от OpenAI, который они используют в своих моделях. Также сделаем все необходимые импорты
import tiktoken
from candle import Tensor
from dataclasses import dataclass
import re # для загрузки весовОпределим конфигурацию нашей модели с помощью dataclasses
@dataclass
class GPTConfig:
block_size: int = 1024 # размер окна вниманиә
vocab_size: int = 50257 # размер словаря BPE
n_layer: int = 12 # количество слоёв
n_head: int = 12 # количество голов в механизме внимания
n_embd: int = 768 # размерность вектора эмбеддинговОпределим линейный слой модели
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
self.gelu = nn.GeLU()
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
def forward(self, x):
x = self.c_fc(x)
x = self.gelu(x)
return self.c_proj(x)Обратите внимание тут candle.nn, а не torch.nn.
С линейными слоями мы уже знакомы, а GeLU(), это что такое?
Исследователи из OpenAI, выяснили, что в их моделях такая функция работает лучше всего, посмотрим на её реализацию
class GeLU:
def __init__(self):
Parameter([self,[]])
def __call__(self, x):
return Tensor.gelu(x)
#return 0.5 * x * (1 + Tensor.tanh(0.79788456 * (x + 0.044715 * (x ** 3))))Я закомментировал вторую строчку и вместо этого использую встроенный метод Tensor.gelu(), но никто не запрещает и честно считать по формуле. Заглянем в этот метод
from scipy.stats import norm
@classmethod
def gelu(cls, x):
return x * norm.cdf(x.value) # формула из картинкиТеперь в линейном слое нам все понятно, идём дальше!
class CausalSelfAttetion(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=True)
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=True)
def forward(self, x):
B, T, C = x.shape
qkv = self.c_attn(x)
q, k, v = Tensor.split(qkv, 3, axis=2)
n_head = self.config.n_head
q = q.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)
k = k.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)
v = v.reshape(B, T, n_head, C // n_head).transpose(0, 2, 1, 3)
att = q @ k.transpose(0, -3, -1, -2) * (k.shape[-1] ** -0.5)
lg = att.local_gradients
att = Tensor.tril(att)
att[att == 0] = float('-inf')
att.local_gradients = lg
probs = Tensor.softmax(att, axis=-1)
y = probs @ v
y = y.transpose(0, 2, 1, 3).reshape(B, T, C)
y = self.c_proj(y)
return yСлой внимания, я предполагаю, что уже видели подобною реализацию, поэтому много уделять внимания не буду. Видим, незнакомые нам методы Tensor.split(), Tensor.tril, они делают тоже самое, что и их аналоги в pytorch
@classmethod
def tril(cls, input, diagonal=0):
value = np.tril(input.value, k=diagonal)
local_gradients = (
('triu', input, lambda x: x * np.tril(np.ones_like(input.value), k=diagonal)),
)
return cls(value, local_gradients=local_gradients)
@classmethod
def split(cls, array, split_size_or_sections, axis=0):
value = np.split(array.value, split_size_or_sections, axis=axis)
return cls(value, requires_grad=False)Обратите внимание, я не определил вычисление градиента для второй операции. Значит, либо мне придется его определить либо я просто не смогу обучать модели!
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd)
self.attn = CausalSelfAttetion(config)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return xТут я использую готовый слой candle.nn.LayerNorm(), заглянем в него!
class LayerNorm:
def __init__(self, dim, eps=1e-5):
super().__init__()
self.param = None
self.gamma = Tensor.ones(dim, requires_grad=True)
self.beta = Tensor.zeros(dim, requires_grad=True)
self.eps = eps
self.all_layers = [self.gamma, self.beta]
self.grad = None
def __call__(self, x):
xmean = Tensor.mean(x, axis=2, keepdims=True)
xstd = Tensor.std(x, axis=2, keepdims=True)
x = (x - xmean) / (xstd + self.eps)
return self.gamma * x + self.betaВ целом ничего сложного, если вы идейно знакомы с методами нормализации. Тут я использую Tensor.mean(), Tensor.std(). Посмотрим на их реализацию
@classmethod
def mean(cls, array, axis=None, keepdims=False):
if axis == None:
return Tensor.sum(array, axis=None, keepdims=keepdims) / np.size(array.value)
else:
delimeter = 1
if not isinstance(axis, int):
for ax in axis:
delimeter = delimeter * array.shape[ax]
else:
delimeter = array.shape[axis]
return Tensor.sum(array, axis=axis, keepdims=keepdims) / delimeter
@classmethod
def std(cls, array, axis=None, keepdims=False):
if axis == None or axis == 0:
mean = Tensor.mean(array, axis=axis, keepdims=False)
sub = array - mean
squared = sub ** 2
scaled_sum = Tensor.mean(squared, axis=axis, keepdims=keepdims)
std = Tensor.sqrt(scaled_sum)
return std
elif axis >= 1:
mean = Tensor.mean(array, axis=axis, keepdims=True)
sub = array - mean
squared = sub ** 2
scaled_sum = Tensor.mean(squared, axis=axis, keepdims=keepdims)
std = Tensor.sqrt(scaled_sum)
out = Tensor(std, local_gradients=None)
out.local_gradients = (('std', std, lambda x: x * array.shape[axis] / (array.shape[axis] - 1)),) # array.shape[axis] / (array.shape[axis] - 1) additional multiplier due to dissimilarity
return out
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.wte = nn.Embedding(config.vocab_size, config.n_embd)
self.wpe = nn.Embedding(1024, config.n_embd)
self.h = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
self.ln_f = nn.LayerNorm(config.n_embd)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
def forward(self, x):
B, T = x.shape
assert T <= self.config.block_size, f"Cannot forward the sequence of length {T}, block_size is smaller"
pos = Tensor.arange(T).reshape(1, -1)
pos_embd = self.wpe(pos)
tok_embd = self.wte(x)
x = pos_embd + tok_embd
for block in self.h:
x = block(x)
x = self.ln_f(x)
logits = self.lm_head(x)
return logits
Разберемся как работают candle.nn.Embedding() и candle.nn.ModuleList.
class Embedding:
def __init__(self, num_emb, emb_dim):
from scipy.sparse import csr_matrix
self.w = Tensor.randn((num_emb, emb_dim), requires_grad=True)
self.w.local_gradients = None
self.num_embd = num_emb
self.emb_dim = emb_dim
self.param = None
self.grad = None
self.all_layers = [self.w]
def __call__(self, x):
if self.param:
global Parameter
Parameter = self.param
def multiply_by_locgrad(path_value):
temp = np.zeros_like(self.w.value)
np.add.at(np.zeros_like(self.w.value), x.value, path_value)
return temp
x.value = x.value.astype(int)
local_gradients = (('embd', self.w, multiply_by_locgrad),)
return Tensor(self.w.value[x.value], local_gradients=local_gradients)В целом ничего сложного, просто ожидаем на входе тензор из целых значений и рассматриваем эти значения как индексы для матрицы хранящей эмбеддинги.
class ModuleList:
def __init__(self, layers):
self.layers = layers
self.index = 0
def __call__(self, x):
for layer in self.layers:
x = layer(x)
return x
def __len__(self):
return len(self.layers)
def __iter__(self):
self.index = 0
return self
def __next__(self):
if self.index < len(self.layers):
result = self.layers[self.index]
self.index += 1
return result
else:
raise StopIteration
def __getitem__(self, index):
return self.layers[index]Оказывается это просто итератор!
Note! Для корректной работы, нам нужно определить метод Module.__setattr__(), иначе слои которые находятся внутри ModuleList, просто не будут видны нашей глобальной переменной Parameter, я не буду на этом останавливаться, при желании можно заглянуть в код и разобраться!
class GPT(nn.Module):
@classmethod
def from_pretrained(cls, model_type):
assert model_type in ('gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl')
config_args = {
'gpt2': dict(n_layer=12, n_head=12, n_embd=768),
'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024),
'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280),
'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600)
}[model_type]
config_args['vocab_size'] = 50257
config_args['block_size'] = 1024
config = GPTConfig(**config_args)
model = GPT(config)
model = GPT.get_params(model, model_type)
return modelЗдесь мы определяем модель из библиотеки hf.transformers, это самый простой путь для доступа к весам GPT2.
@staticmethod
def get_params(model, model_type):
from transformers import GPT2LMHeadModel, logging
logging.set_verbosity_error()
model_hf = GPT2LMHeadModel.from_pretrained(model_type)
sd_hf = model_hf.state_dict()
...Это часть кода, отвечающего за перенос весов с модели hugginface на нашу модель. Там ничего умного, кроме аккуратной работы.
Генерация
model = GPT.from_pretrained(model_type=model_type)
max_length = 100
number_of_examples = 3
starting_sentence = "Hello, I'm a language model,"
enc = tiktoken.get_encoding('gpt2')
tokens = enc.encode(starting_sentence)
tokens = Tensor(tokens, dtype=int)
tokens = Tensor.unsqueeze(tokens, 0)
tokens = Tensor.repeat(tokens, number_of_examples, 0)
topk = args.topk # responsible for "creativity" or "adequacy"
x = tokens
topk = 50 # responsible for "creativity" or "adequacy"
for i in tqdm(range(max_length)):
logits = model(x)
logits = logits[:, -1, :]
probs = Tensor.softmax(logits, axis=-1)
topk_indices, topk_probs = Tensor.topk(probs, topk)
new_token = Tensor.multinomial_from_array(topk_indices, topk_probs, num_samples=1).reshape(-1, 1)
x = Tensor.cat([x, new_token], axis=1)
for sample in x:
sample = sample.value.astype(int).tolist()
print(enc.decode(sample), end='\n')
print('---------------')
Hello, I'm a language model, not a coding model, but a compiler model.
I am thinking of what it means to be a developer and a language model
and to be able to write in some simple yet expressive way code that makes it
possible to use it with our software in any meaningful way to make
applications more readable, better, more readable.
I can think of more words to talk about this.
In this post, I'm trying to explain some of the main differences
---------------
Hello, I'm a language model, based on my personal experience with code and development using C#.
This tutorial will help you create my own virtual language that you can use to implement things like an HTML page in your own app.
Now, what if I were going to learn how to program a website in an Objective-C program and create a web app out of it,
but didn't know how to do that and want to do some extra work to try it out? First I
---------------
Hello, I'm a language model, I'm my own language. Well, a language model has two parts:
a semantic construct and a structural construct.
To understand the relationship with the semantic construct, let's see how all of the variables on a graph come together.
Now I know that some of the variables on a graph represent numbers. But when you see the graphs, they're graphs.
In other words, you see the connections and they are networks. We have an exampleВы можете поиграться с этой моделью на каггле или скачав к себе из репозитория smolGPT
Вот и подошла к концу моя мини-серия по созданию собственной библиотеки на NumPy.
Благодарю за внимание!
Надеюсь вы нашли для себя этот цикл статей познавательным!
Если вам понравилось, пожалуйста, поделитесь ими, поставьте upvote и поставьте звёздочки моим реализациям на гитхабе.
Первая версия библиотеки
Вторая версия библиотеки
GPT-2 на этой библиотеке
