[SOLVED] Trouble using pre-trained model from niftynet

Trouble using pre-trained model from niftynet

I want to do a 2 class segmentation using dense_vnet model available on niftynet which originally does a 9 class segmentation

I tried to retrain only the last layer by making changes in config file according to this suggestion: HOw to fine tune niftynet pre trained model for custom data

vars_to_restore = ^((?!DenseVNet\/(skip_conv|fin_conv)).)*$

num_classes = 2

error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Assign requires shapes of both tensors to match. lhs shape= [2] rhs shape= [9] [[{{node save/Assign_8}} = Assign[T=DT_FLOAT, class=["loc:@DenseVNet/conv/conv/b"], use_locking=true, validate_shape=true, device="/job:localhost/replica:0/task:0/device:CPU:0"](DenseVNet/conv/conv/b, save/RestoreV2:8)]]

Solution

It looks like you have restored too many layers, some of them are still designed to classify to 9 classes. Inspect the architecture and exclude restore for all layers which are designed to classify into 9 classes.