please write the code about how to do the supervised learning in pytorch
时间: 2024-05-01 18:17:28 浏览: 112
Sure, here is an example code for supervised learning in PyTorch:
```python
import torch
import torch.nn as nn
import torch.optim as optim
# Define the neural network
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(in_features=784, out_features=256)
self.fc2 = nn.Linear(in_features=256, out_features=10)
def forward(self, x):
x = torch.flatten(x, start_dim=1)
x = nn.functional.relu(self.fc1(x))
x = self.fc2(x)
return x
# Load the data
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64)
# Instantiate the model, loss function, and optimizer
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# Train the model
for epoch in range(10):
total_loss = 0
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
total_loss += loss.item()
print("Epoch {}: Loss = {:.4f}".format(epoch+1, total_loss / len(train_loader)))
```
In this code, we define a simple neural network model with two fully connected layers for classification. We then load the training data using a PyTorch DataLoader object, instantiate the model, loss function (cross-entropy), and optimizer (stochastic gradient descent), and train the model using a loop over the training data. At each iteration, we compute the output of the model given the input data, calculate the cross-entropy loss between the output and the true labels, compute its gradients, and update the model weights using the optimizer. This process is repeated for a fixed number of epochs, and the training loss is printed for each epoch.
Please note that this is just an example code and needs to be adapted to your specific use case.
阅读全文