我一直在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()
1条答案
按热度按时间fv2wmkja1#
您正在使用HuggingFace Transformers,您可以用途:
保存模型后,文件夹将包含
pytorch_model.bin
沿着配置JSON。