如何为这个 python 文件使用 GPU

Question

我有这个 python 文件，我试图从头开始训练 GPT2 模型。同样，我想使用 gpu 来获得更快的加速，但我无法这样做。将不胜感激

我的python代码如下

PS : 我在 AWS Sagemaker 上运行ning 这段代码，所以我想使用他们的 gpu 加速。

我参考了这个link

import tensorflow as tf
from transformers import GPT2Config, TFGPT2LMHeadModel, GPT2Tokenizer
from pathlib import Path
save_path = 'tokenized_data_clean'
tokenizer = GPT2Tokenizer.from_pretrained(save_path)
from transformers import WEIGHTS_NAME, CONFIG_NAME
tokenizer.add_special_tokens({
    "eos_token":"</s>",
    "bos_token":"<s>",
    "unk_token":"<unk>",
    "pad_token":"<pad>",
    "mask_token":"<mask>"
})

config = GPT2Config(
    vocab_size = tokenizer.vocab_size,
    bos_token_id = tokenizer.bos_token_id,
    eos_token_id = tokenizer.eos_token_id
)

model = TFGPT2LMHeadModel(config)
paths = ['text_data/clean_data_file.txt']
single_string = ''


for filename in paths:
    with open(filename, "r", encoding = 'utf-8') as f:
        x = f.read()

        single_string+=x+tokenizer.eos_token

string_tokenized = tokenizer.encode(single_string)
examples = []
block_size = 100
BATCH_SIZE = 12
BUFFER_SIZE = 1000
for i in range(0, len(string_tokenized) - block_size + 1, block_size):
  examples.append(string_tokenized[i:i + block_size])
inputs, labels = [], []
for ex in examples:
  inputs.append(ex[:-1])
  labels.append(ex[1:])
dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)



# defining our optimizer
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
# definining our loss function
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# defining our metric which we want to observe
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
# compiling the model
model.compile(optimizer=optimizer, loss=[loss, *[None] * model.config.n_layer], metrics=[metric])


num_epoch = 100
history = model.fit(dataset, epochs=num_epoch)

text = "I think"
# encoding the input text
input_ids = tokenizer.encode(text, return_tensors='tf')
# getting out output
beam_output = model.generate(
  input_ids,
  max_length = 50,
  num_beams = 5,
  temperature = 0.7,
  no_repeat_ngram_size=2,
  num_return_sequences=5
)
print(beam_output)

import os
output_dir = './model_bn_custom_100/'
# creating directory if it is not present
if not os.path.exists(output_dir):
  os.mkdir(output_dir)
model_to_save = model.module if hasattr(model, 'module') else model
output_model_file = os.path.join(output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(output_dir, CONFIG_NAME)
# save model and model configs
model.save_pretrained(output_dir)
model_to_save.config.to_json_file(output_config_file)
# save tokenizer
tokenizer.save_pretrained(output_dir)

代码运行良好，我只是希望它使用 gpu 加速运行更快。

Answer 1

在 AWS SageMaker 中托管笔记本会话时，您需要激活 GPU 运行时。该代码将自动处理 GPU 资源的利用。

查看您共享的 link - 它没有任何自定义配置来手动指定 GPU 资源。

如果它由您用来训练网络的框架自动处理，那么在活动的 GPU 会话中，它会在训练时自动分配 GPU 资源。

如何为这个 python 文件使用 GPU

How to use GPU for this python file

python

machine-learning

keras

tensorflow