File size: 1,696 Bytes
2c07569 |
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 |
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')
|