def test_mean_iou(self):
criterion = MeanIoU()
x = torch.randn(3, 3, 3, 3)
_, index = torch.max(x, dim=0, keepdim=True)
# create target tensor
target = torch.zeros_like(x, dtype=torch.long).scatter_(0, index, 1)
pred = torch.zeros_like(target, dtype=torch.float)
mask = target == 1
# create prediction tensor
pred[mask] = torch.rand(1)
# make sure the dimensions are right
target = torch.unsqueeze(target, dim=0)
pred = torch.unsqueeze(pred, dim=0)
assert criterion(pred, target) == 1
def test_mean_iou_one_channel(self):
criterion = MeanIoU()
pred = torch.rand(1, 1, 3, 3, 3)
target = pred > 0.5
target = target.long()
assert criterion(pred, target) == 1
def test_mean_iou_simple(self):
results = _compute_criterion(MeanIoU())
# check that all of the coefficients belong to [0, 1]
results = np.array(results)
assert np.all(results > 0)
assert np.all(results < 1)
elif LOSS_CRITERION == 'BCELoss':
print("Using BCELoss")
loss_criterion = nn.BCEWithLogitsLoss()
elif LOSS_CRITERION == 'CrossEntropyLoss':
print("Using CrossEntropyLoss")
loss_criterion = nn.CrossEntropyLoss()
elif LOSS_CRITERION == 'GeneralizedDiceLoss':
print("Using GeneralizedDiceLoss")
loss_criterion = GeneralizedDiceLoss()
else:
print("No loss criterion specified, exiting")
sys.exit(1)
# Get evaluation metric
if EVAL_METRIC == "MeanIoU":
print("Using MeanIoU")
eval_criterion = MeanIoU()
elif EVAL_METRIC == "GenericAveragePrecision":
print("Using GenericAveragePrecision")
eval_criterion = GenericAveragePrecision()
else:
print("No evaluation metric specified, exiting")
sys.exit(1)
# Create model and optimizer
os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_DEVICE
unet = create_unet_on_device(DEVICE_NUM, MODEL_DICT)
optimizer = torch.optim.AdamW(unet.parameters(), lr=STARTING_LR)
train_subject = torchio.Subject(
data=torchio.Image(tensor=torch.from_numpy(train_data),
label=torchio.INTENSITY),
label=torchio.Image(tensor=torch.from_numpy(train_seg),