Background（背景信息）

Our team tried to implement the DETR model, which uses the Hungarian algorithm for selecting boxes for assignments during the training.

This Hungarian algorithm uses 3 nested while loops, and operates the tensors, which shape during computation dynamically changes. It makes it impossible to use it in GRAPH_MODE and works extremely slow in PYNATIVE_MODE.

Currently, the only solution is to compute losses manually using NumPy and then assigning the manually computed gradients as the sense parameter for the GradientOp operator.

The link to the DETR implementation: !2266:Models: DETR

The current solution can be described as the following:

net = MyNetwork(...)
one_step_cell = CustomTrainOneStepCellWithSense(net)

# ...
# Training cycle
input_data, gt = next(dataset_iterator)

# Perform a forward step to calculate the network outputs and the loss.
pred = net(input_data)

# Manually compute the loss and provide the back-propagated gradients for the loss.
loss, loss_grad = my_loss_function(pred, gt)

# Update the sense parameter of one-step-cell so it starts 
# its back-propagation starting from those values
one_step_cell.sense_param.set_data(loss_grad)

# Performing one more forward step, calculating gradients and the weights update.
one_step_cell(input_data)

The additional forward pass takes additional time, so this solution is a bit slower than the original PyTorch model.

Benefit / Necessity （价值/作用）

The support for the arbitrary functions will allow wider class of models, supported by MindSpore.

Design（设计方案）

The design of a possible solution:

Adding the PythonCell class, where a user can define forward and backward propagation logic, or at least the CustomLossCell operator which accepts the tensor, and returns the loss value and the computed gradients. The internal computation of the CustomLossCell should support arbitrary functions.