File size: 3,370 Bytes
0ab2514 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import logging
import os
import psutil
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
from datasets import load_dataset
import torch
# Configurar logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Usando dispositivo: {device}")
def log_gpu_memory():
if torch.cuda.is_available():
gpu_memory = torch.cuda.memory_allocated() / 1024 / 1024 # em MB
logger.info(f"Uso de memória GPU: {gpu_memory:.2f} MB")
else:
logger.info("GPU não disponível")
# Função para monitorar o uso de RAM
#def log_memory_usage():
# process = psutil.Process(os.getpid())
# mem = process.memory_info().rss / 1024 / 1024 # em MB
# logger.info(f"Uso de memória atual: {mem:.2f} MB")
# Carregue o modelo e o tokenizer
logger.info("Carregando modelo e tokenizer...")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
tokenizer = AutoTokenizer.from_pretrained("t5-small")
log_gpu_memory()
# Prepare seu dataset
logger.info("Carregando dataset...")
dataset = load_dataset('csv', data_files='../extract-csv/social_keywords.csv')
log_gpu_memory()
# Função para preparar os dados
def preprocess_function(examples):
inputs = examples['keyword']
targets = examples['context']
model_inputs = tokenizer(inputs, max_length=128, truncation=True, padding='max_length')
# Tokenize targets
with tokenizer.as_target_tokenizer():
labels = tokenizer(targets, max_length=128, truncation=True, padding='max_length')
model_inputs['labels'] = labels['input_ids']
return model_inputs
# Aplicar pré-processamento
logger.info("Pré-processando dataset...")
train_dataset = dataset['train'].map(preprocess_function, batched=True, remove_columns=dataset['train'].column_names)
log_gpu_memory()
print("Amostra de dados processados:")
print(train_dataset[:5])
# Configure os argumentos de treinamento
logger.info("Configurando argumentos de treinamento...")
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
per_device_train_batch_size=16,
save_steps=10_000,
save_total_limit=2,
logging_dir='./logs',
logging_steps=100,
)
# Classe personalizada de Trainer para monitorar o uso de memória
class MemoryMonitorTrainer(Trainer):
def training_step(self, model, inputs):
loss = super().training_step(model, inputs)
if self.state.global_step % 100 == 0: # Log a cada 100 steps
log_gpu_memory()
return loss
# Crie um data collator
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding='max_length', max_length=128)
# Inicie o treinamento
logger.info("Iniciando treinamento...")
trainer = MemoryMonitorTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
tokenizer=tokenizer,
data_collator=data_collator,
)
trainer.train()
logger.info("Treinamento concluído!")
log_gpu_memory()
# Salvar o modelo treinado
output_dir = "./meu_modelo_treinado"
logger.info(f"Salvando o modelo treinado em {output_dir}")
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
log_gpu_memory()
logger.info("Processo finalizado!") |