d795f811eb13a20e18aa30c435cef0ad99e830ea,examples/plot_unet_tutorial.py,,,#,127

Before Change

    return trainer


trainer = train_model(model=model_a, loaders=[train_loader, validate_loader], save_dir="model_a")



////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Prediction
// ----------------------------
// The trainer contains the trained model and we can do predictions.
// We use :code:`unwrap` to convert the results to numpy arrays.
trainer.eval_mode()
from inferno.utils.torch_utils import unwrap


for image, target in test_loader:

    // transfer image to gpu
    image = image.cuda() if USE_CUDA else image

    // get batch size from image
    batch_size = image.size()[0]
    
    prediction = trainer.apply_model(image)

    image = unwrap(image,      as_numpy=True, to_cpu=True)
    prediction = unwrap(prediction, as_numpy=True, to_cpu=True)


    fig = pylab.figure()

After Change

// or use a some convolutions to increase the number of channels 
// before the the UNet.
from inferno.extensions.layers import ConvReLU2D
model_a = torch.nn.Sequential(
    ConvReLU2D(in_channels=image_channels, out_channels=5, kernel_size=3),
    ResBlockUNet(dim=2, in_channels=5, out_channels=pred_channels,  activated=False,
        res_block_kwargs=dict(batchnorm=True,size=2)) ,
    RemoveSingletonDimension(dim=1)
    // torch.nn.Sigmoid()
)

Instances