search

neuronets / nobrainer

A framework for developing neural network models for 3D image processing.
Other
158 stars 45 forks source link

Add test_unetr() in model_tests.py #301

Open hvgazula opened 8 months ago

hvgazula commented 8 months ago

related to https://github.com/neuronets/nobrainer/issues/302

hvgazula commented 8 months ago

The test works fine on a GPU but not on the github runner (CPU and 16 GB).

Here's a working test-

    from nobrainer.models import unetr
    import numpy as np

    model = unetr()
    model.summary()

    X_train = np.random.rand(1, 96, 96, 96, 3)
    y_train = np.random.rand(1, 96, 96, 96, 1)

    model.compile(loss="binary_crossentropy", optimizer="adam")
    model.fit(X_train, y_train)
    print(model.predict(X_train).shape)

The original code (on MONAI) is also locked to the specific input size, the number of transformer blocks, and other details. I will leave this issue open. This could be a good hackathon issue.