I want to calculate IoU for my segmentation task this is the model used
class UNetModel(nn.Module):
def __init__(self, input_channels, nclasses, threshold=.6):
super().__init__()
self.threshold = threshold
# go down
self.conv1 = conv_bn_leru(input_channels,64)
self.conv2 = conv_bn_leru(64, 128)
self.conv3 = conv_bn_leru(128, 256)
self.conv4 = conv_bn_leru(256, 512)
self.conv5 = conv_bn_leru(512, 1024)
self.down_pooling = nn.MaxPool2d(2)
# go up
self.up_pool6 = up_pooling(1024, 512)
self.conv6 = conv_bn_leru(1024, 512)
self.up_pool7 = up_pooling(512, 256)
self.conv7 = conv_bn_leru(512, 256)
self.up_pool8 = up_pooling(256, 128)
self.conv8 = conv_bn_leru(256, 128)
self.up_pool9 = up_pooling(128, 64)
self.conv9 = conv_bn_leru(128, 64)
self.conv10 = nn.Conv2d(64, nclasses, 1)
# test weight init
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
# go down
x1 = self.conv1(x)
p1 = self.down_pooling(x1)
x2 = self.conv2(p1)
p2 = self.down_pooling(x2)
x3 = self.conv3(p2)
p3 = self.down_pooling(x3)
x4 = self.conv4(p3)
p4 = self.down_pooling(x4)
x5 = self.conv5(p4)
# go up
p6 = self.up_pool6(x5)
x6 = torch.cat([p6, x4], dim=1)
x6 = self.conv6(x6)
p7 = self.up_pool7(x6)
x7 = torch.cat([p7, x3], dim=1)
x7 = self.conv7(x7)
p8 = self.up_pool8(x7)
x8 = torch.cat([p8, x2], dim=1)
x8 = self.conv8(x8)
p9 = self.up_pool9(x8)
x9 = torch.cat([p9, x1], dim=1)
x9 = self.conv9(x9)
output = self.conv10(x9)
output = torch.sigmoid(output)
return output
this is the function that computes IoU:
def iou_pytorch(outputs, labels):
outputs = outputs.squeeze(1) # BATCH x 1 x H x W => BATCH x H x W
intersection = (outputs & labels).float().sum((1, 2))
union = (outputs | labels).float().sum((1, 2))
iou = (intersection + SMOOTH) / (union + SMOOTH) # We smooth our devision to avoid 0/0
thresholded = torch.clamp(20 * (iou - 0.5), 0, 10).ceil() / 10
return thresholded.mean()
this is the part where I compute the IoU
def test(model, validation_loader, criterion, device):
"""evaluate the current state of the model
Args:
model (torch.nn.Module): the model you want to to train
validation_loader (torch.utils.data.DataLoader): validation-set dataloader
criterion (callable): the loss function used to evaluate the model
device (torch.device): device used for calculations 'CPU' or 'GPU'
"""
with torch.no_grad():
model.eval()
test_loss = 0
iou = 0
for idx, data in enumerate(validation_loader):
mask, image = data['mask'].to(device), data['image'].to(device)
output = model(image)
# sum up batch loss
test_loss += criterion(output, mask).item()
output = (output > 0.7).type(torch.cuda.DoubleTensor)
iou += iou_pytorch(output, mask.type(torch.cuda.DoubleTensor))
iou /= len(validation_loader.dataset)
test_loss /= len(validation_loader.dataset)
return test_loss, iou
But RuntimeError: cbitand is only supported for integer type tensors Raises.