欢迎访问宙启技术站
智能推送

使用tensorboard_logger在Python中展示深度学习模型的模型权重变化

发布时间:2024-01-09 09:29:50

要使用tensorboard_logger在Python中展示深度学习模型的模型权重变化,首先我们需要安装tensorboard_logger库。可以使用pip命令进行安装:

pip install tensorboard_logger

完成安装之后,我们可以使用以下例子来展示权重变化:

import tensorboard_logger as tb_logger
import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的神经网络模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 创建一个实例
model = Model()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 创建TensorBoard Logger对象
logger = tb_logger.Logger(logdir='./logs', flush_secs=2)

# 训练模型
for epoch in range(100):
    inputs = torch.randn(32, 10)  # 随机生成输入数据
    labels = torch.randint(0, 2, (32,))  # 随机生成标签数据

    # 模型前向传播
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    # 模型反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 记录模型权重变化
    for name, params in model.named_parameters():
        logger.log_value(f'Weights/{name}.mean', params.mean(), epoch)
        logger.log_value(f'Weights/{name}.std', params.std(), epoch)

    # 打印中间结果
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.item()}')

# 启动tensorboard服务器来查看结果
# 命令行执行 tensorboard --logdir=./logs

在上述例子中,我们首先定义了一个简单的神经网络模型,并创建了一个TensorBoard Logger对象,将日志保存到"./logs"目录下。然后,我们使用这个Logger对象在每个epoch中记录模型权重的平均值和标准差,通过调用logger.log_value()方法来记录这些数值。

在训练循环中,我们首先生成随机的输入数据和标签数据,然后进行模型的前向传播和反向传播。然后,我们通过遍历模型的所有参数,获取它们的平均值和标准差,并使用logger.log_value()方法记录这些数值。最后,我们通过tensorboard命令启动tensorboard服务器,并指定logdir参数为logs目录的路径,就可以在浏览器中查看模型权重的变化了。

要查看模型权重的变化,可以在浏览器中打开http://localhost:6006,在页面上选择"Scalars"选项卡,然后在搜索框中输入"Weights"即可找到我们刚刚记录的模型权重变化的曲线图。

希望这个例子对您有所帮助!