ICNN

class deepinv.models.ICNN(in_channels=3, num_filters=64, kernel_dim=5, num_layers=10, strong_convexity=0.5, pos_weights=True, device='cpu')

Bases: Module

Convolutional Input Convex Neural Network (ICNN).

The network is built to be convex in its input. The model is fully convolutional and thus can be applied to images of any size.

Based on the implementation from the paper “Data-Driven Mirror Descent with Input-Convex Neural Networks.

Parameters:

• in_channels (int) – Number of input channels.

• num_filters (int) – Number of hidden units.

• kernel_dim – dimension of the convolutional kernels.

• num_layers (int) – Number of layers.

• strong_convexity (float) – Strongly convex parameter.

• pos_weights (bool) – Whether to force positive weights in the forward pass.

• device (str) – Device to use for the model.

forward(x)

Calculate potential function of the ICNN.

Parameters:

x (torch.Tensor) – Input tensor of shape (B, C, H, W).

grad(x)

Calculate the gradient of the potential function.

Parameters:

x (torch.Tensor) – Input tensor of shape (B, C, H, W).