torch使用tensorboard简明备忘录

2021-12-30

tensorboard是让torch可以使用tensorflow的web可视化工具，之前叫tensorboardX。
至于其他的介绍以及为什么需要，可自行百度。

简单的完整代码

# -*- coding: utf8 -*-
#

import math

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('./run')

epochs = 100
data_loader = range(1000)

total_loss = 0
for epoch in range(epochs):
    for i, data in enumerate(data_loader):
        loss = math.sin(i)
        total_loss += loss
        # 打印每一个batch下的总loss
        writer.add_scalar('train_loss', total_loss, epoch * len(data_loader) + i)
    total_loss = 0

终端跑下面这条命令：

1	tensorboard --logdir=./runs

更完整的demo

其中数据集加载没有添加进来，引用代码可参考lstm_sent_polarity。

# Defined in Section 4.6.7

import torch
from torch import nn, optim
from torch.nn import functional as F
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence
from torch.utils import tensorboard
from torch.utils.data import Dataset, DataLoader
# tqdm是一个Python模块，能以进度条的方式显式迭代的进度
from tqdm.auto import tqdm

from utils import load_sentence_polarity

summary = tensorboard.SummaryWriter('./runs')

class LstmDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        return self.data[i]


def collate_fn(examples):
    lengths = torch.tensor([len(ex[0]) for ex in examples])
    inputs = [torch.tensor(ex[0]) for ex in examples]
    targets = torch.tensor([ex[1] for ex in examples], dtype=torch.long)
    # 对batch内的样本进行padding，使其具有相同长度
    inputs = pad_sequence(inputs, batch_first=True)
    return inputs, lengths, targets


class LSTM(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_class):
        super(LSTM, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, bidirectional=True, num_layers=3)
        self.output = nn.Linear(hidden_dim * 6, num_class)

    def forward(self, inputs, lengths):
        embeddings = self.embeddings(inputs)
        x_pack = pack_padded_sequence(embeddings, lengths, batch_first=True, enforce_sorted=False)
        hidden, (hn, cn) = self.lstm(x_pack)
        outputs = self.output(hn.permute(1, 0, 2).reshape(-1, 6 * 256))
        log_probs = F.log_softmax(outputs, dim=-1)
        return log_probs


embedding_dim = 128
hidden_dim = 256
num_class = 2
batch_size = 32
num_epoch = 20

# 加载数据
train_data, test_data, vocab = load_sentence_polarity()
train_dataset = LstmDataset(train_data)
test_dataset = LstmDataset(test_data)
train_data_loader = DataLoader(train_dataset, batch_size=batch_size, collate_fn=collate_fn, shuffle=True)
test_data_loader = DataLoader(test_dataset, batch_size=1, collate_fn=collate_fn, shuffle=False)

# 加载模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LSTM(len(vocab), embedding_dim, hidden_dim, num_class)
model.to(device)  # 将模型加载到GPU中（如果已经正确安装）

# 训练过程
nll_loss = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 使用Adam优化器
scheduler = optim.lr_scheduler.StepLR(optimizer, 15, gamma=0.99)

model.train()
for epoch in range(num_epoch):
    total_loss = 0
    for i, batch in tqdm(enumerate(train_data_loader), desc=f"Training Epoch {epoch}"):
        inputs, lengths, targets = [x.to(device) for x in batch]
        lengths = lengths.cpu()
        log_probs = model(inputs, lengths)
        loss = nll_loss(log_probs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        # 对每一个batch记录loss变化
        summary.add_scalar('train_batch_loss', total_loss, epoch * len(train_data_loader) + i)
    scheduler.step()
    print(f"Loss: {total_loss:.2f}, lr: {scheduler.get_last_lr()[0]}")
    # 对每一个epoch记录loss变化
    summary.add_scalar('train_epoch_loss', total_loss, epoch)
    # 对每一个epoch记录lr变化
    summary.add_scalar('train_epoch_lr', scheduler.get_last_lr()[0], epoch)

# 测试过程
acc = 0
for batch in tqdm(test_data_loader, desc=f"Testing"):
    inputs, lengths, targets = [x.to(device) for x in batch]
    lengths = lengths.cpu()

    with torch.no_grad():
        output = model(inputs, lengths)
        acc += (output.argmax(dim=1) == targets).sum().item()

# 输出在测试集上的准确率
print(f"Acc: {acc / len(test_data_loader):.2f}")