Spaces:
Configuration error
Configuration error
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
# Set the device to GPU if available, otherwise use CPU | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
# Define the model architecture | |
class Net(nn.Module): | |
def __init__(self, input_size, hidden_size, output_size): | |
super(Net, self).__init__() | |
self.hidden = nn.Linear(input_size, hidden_size) | |
self.relu = nn.ReLU() | |
self.output = nn.Linear(hidden_size, output_size) | |
self.sigmoid = nn.Sigmoid() | |
def forward(self, x): | |
hidden = self.hidden(x) | |
relu = self.relu(hidden) | |
output = self.output(relu) | |
output = self.sigmoid(output) | |
return output | |
# Define the hyperparameters | |
learning_rate = 0.01 | |
num_epochs = 1000 | |
# Create an instance of the model | |
model = Net(input_size=2, hidden_size=5, output_size=1) | |
# Convert the model to run on the device | |
model.to(device) | |
# Define the loss function and optimizer | |
criterion = nn.BCELoss() | |
optimizer = optim.SGD(model.parameters(), lr=learning_rate) | |
# Load the training data | |
X_train = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]]).float().to(device) | |
y_train = torch.tensor([[0], [1], [1], [0]]).float().to(device) | |
# Train the model | |
for epoch in range(num_epochs): | |
# Forward pass | |
outputs = model(X_train) | |
loss = criterion(outputs, y_train) | |
# Backward and optimize | |
optimizer.zero_grad() | |
loss.backward() | |
optimizer.step() | |
# Print the loss at every 100th epoch | |
if (epoch + 1) % 100 == 0: | |
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, loss.item())) | |
# Save the trained model | |
torch.save(model.state_dict(), 'trained model.pt') | |