Spaces:
Sleeping
Sleeping
File size: 1,573 Bytes
5f543e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import torch
import numpy as np
from monai.networks.nets import UNet
from monai.networks.layers import Norm
def test_model_forward_pass():
"""Test 3D U-Net forward pass with dummy input"""
model = UNet(
spatial_dims=3,
in_channels=1,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
)
dummy_input = torch.randn(1, 1, 96, 96, 96)
with torch.no_grad():
output = model(dummy_input)
assert output.shape == (1, 3, 96, 96, 96)
def test_output_channels():
"""Test model has correct number of output classes"""
model = UNet(
spatial_dims=3,
in_channels=1,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
)
dummy_input = torch.randn(1, 1, 96, 96, 96)
with torch.no_grad():
output = model(dummy_input)
assert output.shape[1] == 3 # background, liver, tumor
def test_softmax_output():
"""Test softmax probabilities sum to 1"""
model = UNet(
spatial_dims=3,
in_channels=1,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
)
dummy_input = torch.randn(1, 1, 96, 96, 96)
with torch.no_grad():
output = model(dummy_input)
probs = torch.softmax(output, dim=1)
sums = probs.sum(dim=1)
assert torch.allclose(sums, torch.ones_like(sums), atol=1e-5) |