Can't save my finetuned model

Beginners
1

Hi everyone, I have the following code below but it’s not saving the model. Can anyone give me suggestions or tell me why it’s not saving it?

import torch
from torch.utils.data import DataLoader
from torch.optim import AdamW
import pandas as pd
from app.utils import data_utils
from datasets import Dataset, load_dataset
from transformers import AutoTokenizer, DistilBertForSequenceClassification, get_scheduler
from tqdm.auto import tqdm
import evaluate

# Fine-tuning DistilBERT for news sentiment analysis (using Kaggle news sentiment analysis) -> 3 labels
# 0 = Negative
# 1 = Neutral
# 2 = Positive
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=3)

# Other test news sentiment
# ds = load_dataset("sara-nabhani/ML-news-sentiment")

# Tokenize Function
def tokenize_function(data):
    return tokenizer(data["text"], padding="max_length", truncation=True)


# https://www.kaggle.com/datasets/clovisdalmolinvieira/news-sentiment-analysis?resource=download
# dataset = load_dataset("yelp_review_full")
# print(dataset)


# Load in dataset and shape it into pandas dataframe
df = data_utils.load_raw_data("kaggle_news_sentiment_analysis.csv")
df["text"] = df["Title"] + ": " + df["Description"]
df = df[["text", "Sentiment"]]
replacements = {"negative": 0, "neutral": 1, "positive": 2}
df["Sentiment"] = df["Sentiment"].map(replacements).fillna(df["Sentiment"])
df.rename(columns={"Sentiment": "labels"}, inplace=True)
print(df)


# Create train and test datasets
train = df.head(3000)
test = df.tail(500)


# Convert from pandas dataframe to huggingface dataframe
train = Dataset.from_pandas(train)
print(train)
test = Dataset.from_pandas(test)
print(test)


# Tokenize the datasets
tokenized_train = train.map(tokenize_function, batched=True)
tokenized_train = tokenized_train.remove_columns(["text"])
tokenized_train.set_format(type="torch")
print(tokenized_train)

tokenized_test = test.map(tokenize_function, batched=True)
tokenized_test = tokenized_test.remove_columns(["text"])
tokenized_test.set_format(type="torch")
print(tokenized_test)


# Create data loaders
train_dataloader = DataLoader(tokenized_train, shuffle=True, batch_size=10)
test_dataloader = DataLoader(tokenized_test, batch_size=10)


# Initialize optimizer
optimizer = AdamW(model.parameters(), lr=5e-5)


# Initialize learning rate scheduler
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps)


# Specify device to train on
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)


# Training loop
print("STARTING TRAINING LOOP ==============")
progress_bar = tqdm(range(num_training_steps))

model.train()
for epoch in range(num_epochs):
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch) # Pass batch through model and get output
        loss = outputs.loss # Calculate loss from output
        loss.backward() # Backpropagation using losses to calculate gradient

        optimizer.step() # Update parameters
        lr_scheduler.step() #
        optimizer.zero_grad() # Refresh gradient
        progress_bar.update(1) # Update progress bar

# Save model
model.save_pretrained(save_directory="./saved_models/distilbert")
1 Like
2

This may fix the problem. If it doesn’t, there may be a bug in one of the libraries.

# Save model
model.to("cpu") # Added
model.save_pretrained(save_directory="./saved_models/distilbert")
3

I’ll try it out after this, but I was wondering if it was possible to train the model with native pytorch and still save it using HuggingFace’s save_pretained() function. Would that also be another possible issue? I currently don’t get any error messages; it just doesn’t save it.

1 Like
4

Usually, if save_pretrained fails, an error message is displayed in many cases, so it smells like a bug…
Basically, what is saved with HF’s save_pretrained is the torch model’s state_dict and the configuration json, so converting a model trained with native torch to HF format is usually not that difficult.
The only problem is when the state_dict keys change.

5

I found out the problem. I forgot to put return_tensors=“pt”. :smiling_face_with_tear: :sob: :sob:

1 Like
6

This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.