ENSEEIHT/projet-long
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Solved the problem that wandb could not display masks_pred in the multi-category segmentation problem.If we do not modify it, we may encounter the following error message:mode, rawmode = _fromarray_typemap[typekey] KeyError: ((1, 1, 15), '|u1').ps:15 is the number of categories.

Browse source 
Former-commit-id: ff0c2ffa33c945a491e85ef44e89891e9a7c435d
This commit is contained in:
Liunaijiaaa 2021-11-13 16:41:05 +08:00
parent a445c4bdc4
commit 06199357af
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
train.py
View file

@ -130,7 +130,7 @@ def train_net(net,
'images': wandb.Image(images[0].cpu()), 'images': wandb.Image(images[0].cpu()),
'masks': { 'masks': {
'true': wandb.Image(true_masks[0].float().cpu()), 'true': wandb.Image(true_masks[0].float().cpu()),
'pred': wandb.Image(torch.softmax(masks_pred, dim=1)[0].float().cpu()), 'pred': wandb.Image(torch.softmax(masks_pred, dim=1).argmax(dim=1)[0].float().cpu()),
}, },
'step': global_step, 'step': global_step,
'epoch': epoch, 'epoch': epoch,