<jupyter_start><jupyter_text>此笔记本必须放到与Core同目录下<jupyter_code>import os
import torch
import torchvision
import numpy as np
import matplotlib.pyplot as plt

from Core.Function import sample
from Core.Function import set_seed
from Core.CONFIG import Trainer
from Core.Model import GPT_Model
from Core.Function import kmeans
from Core.Datasets import ImageDataset
from Core.CONFIG import GPTConfig
from Core.CONFIG import TrainerConfig<jupyter_output><empty_output><jupyter_text>下载数据<jupyter_code># ===========================下载数据====================================
# 加载数据
root = './'
train_data = torchvision.datasets.CIFAR10(root, train=True, transform=None, target_transform=None, download=True)
test_data  = torchvision.datasets.CIFAR10(root, train=False, transform=None, target_transform=None, download=True)
print(len(train_data), len(test_data))<jupyter_output>Files already downloaded and verified
Files already downloaded and verified
50000 10000<jupyter_text>处理数据<jupyter_code># ================================================================
# 每张图像随机获取 5 个像素并将它们全部堆叠为 rgb 值以获得半百万个随机像素
pluck_rgb = lambda x: torch.from_numpy(np.array(x)).view(32*32, 3)[torch.randperm(32*32)[:5], :]
px = torch.cat([pluck_rgb(x) for x, y in train_data], dim=0).float()
print(px.size())

# ===========================应用K-means进行获取数据离散值=====================================
ncluster = 512
with torch.no_grad():
    C = kmeans(px, ncluster, niter=8)

print(C.size()) # 输出结果
# =============================制作数据集==============================
train_dataset = ImageDataset(train_data, C)                     
#test_dataset = ImageDataset(test_data, C)                        
print(train_dataset[0][0])  # 一个示例图像被展平为整数<jupyter_output>tensor([239, 312, 310,  ..., 201, 231, 376])<jupyter_text>将预训练的模型加载进去（注意模型必须放到```Pre_models```目录之中）<jupyter_code># 设置确定性
set_seed(42)
# ===================================================================
# 训练前的一些GPT模型的配置
# 根据官方的模型，参数为batch_size = 128,Adam lr 0.003，beta = (0.9, 0.95)
# 学习率预热一个 epoch，然后衰减到 0
# 没有使用权重衰减或Droput
# n_layer=24, n_head=8, n_embd=512
# 另外您可以根据自己的设备进行自己配置
mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size,
                  embd_pdrop=0.0, resid_pdrop=0.0, attn_pdrop=0.0,
                  n_layer=10, n_head=4, n_embd=84)

model = GPT_Model(mconf)
#print(model)
# =============================Load Model=====================================
checkpoint = torch.load('model.bin')
model.load_state_dict(checkpoint)<jupyter_output><empty_output><jupyter_text>转换为ONNX模型<jupyter_code># =============================CONFIG =======================================
tokens_per_epoch = len(train_data) * train_dataset.block_size
train_epochs = 1 # todo run a bigger model and longer, this is tiny
# 初始化训练器进行训练
tconf = TrainerConfig(max_epochs=train_epochs, batch_size=3*8, learning_rate=3e-3,
                      betas = (0.9, 0.95), weight_decay=0,
                      lr_decay=True, warmup_tokens=tokens_per_epoch, final_tokens=train_epochs*tokens_per_epoch,
                      num_workers=8)

trainer = Trainer(model = model, train_dataset = train_dataset, test_dataset = None, config = tconf,Save_Model_path='./pa')
# =============================================================================
counts = torch.ones(ncluster) # start counts as 1 not zero, this is called "smoothing"
rp = torch.randperm(len(train_dataset))
nest = 5000 # how many images to use for the estimation
for i in range(nest):
    a, _ = train_dataset[int(rp[i])]
    t = a[0].item() # index of first token in the sequence
    counts[t] += 1
    
prob = counts/counts.sum()
n_samples = 32
start_pixel = np.random.choice(np.arange(C.size(0)), size=(n_samples, 1), replace=True, p=prob)
start_pixel = torch.from_numpy(start_pixel).to(trainer.device)
import torch.onnx 




    # set the model to inference mode 
model.eval() 

    # Let's create a dummy input tensor  
dummy_input = start_pixel#torch.randn(1, input_size, requires_grad=True)  

    # Export the model   
torch.onnx.export(model,
                  # model being run 
                 dummy_input,       # model input (or a tuple for multiple inputs) 
                 "iGPT.onnx",       # where to save the model  
                 export_params=True,  # store the trained parameter weights inside the model file 
                 opset_version=10,    # the ONNX version to export the model to 
                 do_constant_folding=True,  # whether to execute constant folding for optimization 
                 input_names = ['modelInput'],   # the model's input names 
                 output_names = ['modelOutput'], # the model's output names 
                 dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes 
                                'modelOutput' : {0 : 'batch_size'}})

print(" ") 
print('Model has been converted to ONNX')<jupyter_output>Model has been converted to ONNX<jupyter_text>展示您的模型生成的结果<jupyter_code>n_samples = 32
start_pixel = np.random.choice(np.arange(C.size(0)), size=(n_samples, 1), replace=True, p=prob)
start_pixel = torch.from_numpy(start_pixel).to(trainer.device)

pixels = sample(model, start_pixel, 32*32-1, temperature=1.0, sample=True, top_k=100)
# =========================Show the Images===============================
# for visualization we have to invert the permutation used to produce the pixels
iperm = torch.argsort(train_dataset.perm)
ncol = 8
nrow = n_samples // ncol

plt.figure(figsize=(16, 8))
for i in range(n_samples):
    pxi = pixels[i][iperm] # note: undo the encoding permutation
    
    plt.subplot(nrow, ncol, i+1)
    plt.imshow(C[pxi].view(32, 32, 3).numpy().astype(np.uint8))
    plt.axis('off')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.20195838433292534, 'fraction_numerical': 0.023605525441510754, 'mean_word_length': 4.56420233463035, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '206642', 'license': 'MIT', 'n_tokens_mistral': 2073, 'n_tokens_neox': 1983, 'n_words': 403, 'path': 'iGPT/iGPT_pre.ipynb', 'repo_name': 'StarxSky/GPT', 'size': '5719'}