please explain the code of torch.optim.AdamW
时间: 2024-03-19 19:44:22 浏览: 96
Certainly. Here is an example code snippet for using torch.optim.AdamW:
```
import torch
import torch.optim as optim
# Define the model
model = torch.nn.Linear(10, 1)
# Define the optimizer
optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)
# Define the loss function
criterion = torch.nn.MSELoss()
# Train the model
for epoch in range(100):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
```
In this code snippet, we first define a simple linear model with 10 input features and 1 output feature. We then define an instance of the AdamW optimizer by passing in the model's parameters, a learning rate of 0.001, and a weight decay of 0.01.
We also define a mean squared error loss function (MSELoss) to use for training the model. Finally, we train the model for 100 epochs using a loop. In each epoch, we zero out the gradients, compute the outputs of the model, compute the loss, backpropagate the loss to compute the gradients, and update the model's parameters using the optimizer's step method.
This is a simple example of how to use torch.optim.AdamW for training a neural network.
阅读全文