【发布时间】：2022-08-05 15:15:03
【问题描述】：
我正在使用 pytorch 闪电，在每个时期结束时，我从 torchmetrics.ConfusionMatrix 创建一个混淆矩阵（参见下面的代码）。我想将此记录到 Wandb，但 Wandb 混淆矩阵记录器仅接受 y_targets 和 y_predictions。有谁知道如何从混淆矩阵中提取更新的混淆矩阵 y_targets 和 y_predictions，或者以一种可以将其处理为例如 wandb 内的热图的方式将更新后的混淆矩阵提供给 Wandb？
class ClassificationTask(pl.LightningModule):
def __init__(self, model, lr=1e-4, augmentor=augmentor):
super().__init__()
self.model = model
self.lr = lr
self.save_hyperparameters() #not being used at the moment, good to have ther in the future
self.augmentor=augmentor
self.matrix = torchmetrics.ConfusionMatrix(num_classes=9)
self.y_trues=[]
self.y_preds=[]
def training_step(self, batch, batch_idx):
x, y = batch
x=self.augmentor(x)#.to(\'cuda\')
y_pred = self.model(x)
loss = F.cross_entropy(y_pred, y,) #weights=class_weights_tensor
acc = accuracy(y_pred, y)
metrics = {\"train_acc\": acc, \"train_loss\": loss}
self.log_dict(metrics)
return loss
def validation_step(self, batch, batch_idx):
loss, acc = self._shared_eval_step(batch, batch_idx)
metrics = {\"val_acc\": acc, \"val_loss\": loss, }
self.log_dict(metrics)
return metrics
def _shared_eval_step(self, batch, batch_idx):
x, y = batch
y_hat = self.model(x)
loss = F.cross_entropy(y_hat, y)
acc = accuracy(y_hat, y)
self.matrix.update(y_hat,y)
return loss, acc
def validation_epoch_end(self, outputs):
confusion_matrix = self.matrix.compute()
wandb.log({\"my_conf_mat_id\" : confusion_matrix})
def configure_optimizers(self):
return torch.optim.Adam((self.model.parameters()), lr=self.lr)
标签： pytorch confusion-matrix pytorch-lightning wandb