给一个代码实例,在pytorch中使用tensorboard显示训练过程和结果
时间: 2024-02-03 10:15:10 浏览: 27
好的,下面是一个使用 PyTorch 和 TensorBoard 的示例代码:
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
# 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 1)
def forward(self, x):
x = self.fc1(x)
x = torch.relu(x)
x = self.fc2(x)
return x
# 创建模型和优化器
net = Net()
optimizer = optim.SGD(net.parameters(), lr=0.1)
# 创建TensorBoard 编写器
writer = SummaryWriter()
# 训练模型
for epoch in range(100):
# 假设这里有一批数据 input 和 target
input = torch.randn(10)
target = torch.randn(1)
# 前向传播计算损失
output = net(input)
loss = nn.functional.mse_loss(output, target)
# 反向传播更新参数
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 记录训练过程和结果到TensorBoard
writer.add_scalar('Loss/train', loss.item(), epoch)
writer.add_scalar('Accuracy/train', accuracy, epoch)
writer.add_histogram('conv1.weight', net.conv1.weight, epoch)
# 关闭TensorBoard编写器
writer.close()
```
在上面的代码中,我们首先定义了一个包含两个全连接层的神经网络,然后创建了一个 SGD 优化器和一个 TensorBoard 编写器。在训练过程中,我们使用随机生成的数据进行前向传播和反向传播,并使用 `SummaryWriter` 将训练过程和结果记录到 TensorBoard 中。最后,我们关闭编写器并结束训练过程。