How to export a PyTorch model for HuggingFace?

Question

I have been training my custom Image classification model on the PyTorch transformers library to deploy to hugging face however, I cannot figure out how to export the model in the correct format for HuggingFace with its respective config.json file.

I'm new to PyTorch and AI so any help would be greatly appreciated

train.py

from tqdm import tqdm

best_accuracy = 0

# Train the model for a number of epochs
for epoch in range(20):
    # Create a progress bar for this epoch
    pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{20}')
    
    # Loop over each batch of data
    for X_batch, y_batch in pbar:
        # Move the batch of data to the device
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        # Zero the gradients...
        # Define an optimizer...
        
        # Update the progress bar
        pbar.set_postfix({'Loss': loss.item()})
    
    # Evaluate the model on the validation set
    model.eval()
    correct = 0
    total = 0
    val_loss = 0
    
    with torch.no_grad():
        for X_batch, y_batch in test_loader:
            # Move the batch of data to the device
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)
            
            # Compute the model's predictions for this batch of data
            y_pred = model(X_batch)
            
            # Compute the loss
            loss = criterion(y_pred, y_batch)
            val_loss += loss.item()
            
            # Compute the number of correct predictions
            _, predicted = torch.max(y_pred.data, 1)
            total += y_batch.size(0)
            correct += (predicted == y_batch).sum().item()
    
    val_loss /= len(test_loader)
    accuracy = correct / total
    
    print(f'Validation Loss: {val_loss:.4f}, Accuracy: {accuracy:.4f}')
    
    if accuracy > best_accuracy:
        best_accuracy = accuracy
        torch.save(model.state_dict(), 'best_model.pth')
    
    model.train()

Answer 1

You are using HuggingFace Transformers, you can use:

model.save_pretrained("FOLDER_NAME_HERE")

After you saved the model, the folder will contain the pytorch_model.bin along with config JSONs.