era-resnet / create_demo_model.py
Arnab Sinha
Convert to Hugging Face Gradio app
b12618f
Raw
History Blame Contribute Delete
948 Bytes
import torch
import torch.nn as nn
from model import resnet18_cifar
import os
def create_demo_checkpoint():
"""
Create a demo checkpoint for demonstration purposes.
In a real deployment, you would use a properly trained model.
"""
model = resnet18_cifar(num_classes=100, width=64)
# Initialize with random weights (in practice, use trained weights)
checkpoint = {
'epoch': 100,
'model_state': model.state_dict(),
'best_acc1': 75.0, # Example accuracy
'args': {
'width': 64,
'num_classes': 100,
'max_lr': 0.1,
'weight_decay': 5e-4,
}
}
# Save demo checkpoint
os.makedirs('checkpoints', exist_ok=True)
torch.save(checkpoint, 'checkpoints/demo_model.pth')
print("Demo checkpoint created at checkpoints/demo_model.pth")
return checkpoint
if __name__ == "__main__":
create_demo_checkpoint()