如何为HuggingFace导出PyTorch模型?

iyfjxgzm  于 2023-06-06  发布在  其他
关注(0)|答案(1)|浏览(190)

我一直在PyTorch transformers库上训练我的自定义图像分类模型,以部署到拥抱脸,但是,我无法弄清楚如何使用其各自的config.json文件以正确的格式导出HuggingFace的模型。
我是PyTorch和AI的新手,因此任何帮助都将不胜感激
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()
fv2wmkja

fv2wmkja1#

您正在使用HuggingFace Transformers,您可以用途:

model.save_pretrained("FOLDER_NAME_HERE")

保存模型后,文件夹将包含pytorch_model.bin沿着配置JSON。

相关问题