File size: 2,782 Bytes
6dc5627 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import json
import os
# Step 1: Define Your Dataset Class
class CustomDataset(Dataset):
def __init__(self, texts, labels):
self.texts = texts
self.labels = labels
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
return self.texts[idx], self.labels[idx]
# Step 2: Define Your Model Class
class LSTMModel(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(LSTMModel, self).__init__()
self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x):
lstm_out, _ = self.lstm(x)
out = self.fc(lstm_out[:, -1, :]) # Get the last time step output
return out
# Step 3: Initialize Hyperparameters and Model
input_size = 100 # Example input size (e.g., embedding size)
hidden_size = 64 # Number of LSTM units
output_size = 10 # Number of output classes
num_epochs = 5
learning_rate = 0.001
# Initialize the model
model = LSTMModel(input_size, hidden_size, output_size)
# Step 4: Set Up Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# Step 5: Sample Data (You would replace this with your actual data)
texts = torch.randn(100, 10, input_size) # 100 samples, sequence length of 10
labels = torch.randint(0, output_size, (100,)) # 100 random labels
# Create a DataLoader
dataset = CustomDataset(texts, labels)
data_loader = DataLoader(dataset, batch_size=16, shuffle=True)
# Step 6: Training Loop
for epoch in range(num_epochs):
for inputs, targets in data_loader:
# Forward pass
outputs = model(inputs)
loss = criterion(outputs, targets)
# Backward pass and optimization
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
# Step 7: Save the Model
model_save_path = "model" # Change this to your desired path
os.makedirs(model_save_path, exist_ok=True) # Create the directory if it doesn't exist
# Save the model weights as pytorch_model.bin
torch.save(model.state_dict(), os.path.join(model_save_path, "pytorch_model.bin"))
# Step 8: Create and Save the Configuration File
config = {
"input_size": input_size,
"hidden_size": hidden_size,
"output_size": output_size,
"num_layers": 1, # Add more parameters as needed
"dropout": 0.2
}
# Save the configuration to a JSON file
with open(os.path.join(model_save_path, "config.json"), "w") as f:
json.dump(config, f)
print("Model and configuration saved successfully!")
|