def forward(ctx, input, proj):
batchsize = input.shape[0]
ctx.save_for_backward(input)
if (input.shape[1] != 1):
raise NotImplementedError
sp = Config.getScanParam()
out = input.clone()
for i in range(batchsize):
cupy_input = fromDlpack(to_dlpack(out[i, 0, :, :]))
cupy_proj = fromDlpack(to_dlpack(proj[i, 0, :, :]))
CTOperator.SART2D(cupy_proj, sp, sp['order'], cupy_input)
return out
def update_W(W, a_last, z, u, rho):
size = dist.get_world_size()
rank = dist.get_rank()
# convert to pytorch data
#update W
temp1 = z + u / rho
temp1 = from_dlpack(toDlpack(temp1))
a_last = from_dlpack(toDlpack(a_last))
data1 = torch.mm(temp1, torch.t(a_last))
data2 = torch.mm(a_last, torch.t(a_last))
data = torch.cat((data1, data2), 0)
# data = comm.reduce(data, op=MPI.SUM, root=0)
dist.reduce(data, dst=0, op=dist.ReduceOp.SUM)
if rank == 0:
middle_pos = data1.shape[0]
data1 = data[0:middle_pos]
data2 = data[middle_pos:]
inverse_data = torch.pinverse(data2)
W = torch.mm(data1, inverse_data)
else:
W = from_dlpack(toDlpack(W))
# W = None
dist.broadcast(W, src=0)
# convert to cupy data
W = fromDlpack(to_dlpack(W))
return W
def backward(ctx, grad_output):
batchsize = grad_output.shape[0]
# input = ctx.saved_varibles
if (grad_output.shape[1] != 1):
raise NotImplementedError
sp = Config.getScanParam()
grad = grad_output.clone()
for i in range(batchsize):
cupy_grad = fromDlpack(to_dlpack(grad[i, 0, :, :]))
CTOperator.SART2DBackWard(cupy_grad, sp['order'], sp)
return grad, None
show_num = 10 # 20个batch显示一次
if (i + 1) % show_num == 0:
print('\n ---- batch: %03d of epoch: %03d ----' % (i + 1, epoch + 1))
print('cla_loss: %f, reg_loss: %f, box_loss: %f, aux_seg_loss: %f, aux_offset_loss: %f' \
% (batch_cla_loss / show_num, batch_reg_loss / show_num, batch_box_loss / show_num, \
batch_aux_seg_loss / show_num, batch_aux_offset_loss / show_num))
print('accuracy: %f' % (batch_correct / float(batch_num)))
print('true accuracy: %f' % (batch_true_correct / show_num))
# vis.plot('cla_loss', (batch_cla_loss / show_num).item())
# vis.plot('reg_loss', (batch_reg_loss / show_num).item())
# vis.plot('box_loss', (batch_box_loss / show_num).item())
# vis.plot('aux_seg_loss', (batch_aux_seg_loss / show_num).item())
# vis.plot('aux_offset_loss', (batch_aux_offset_loss / show_num).item())
# vis.plot('accuracy', (batch_correct / float(batch_num)))
# vis.plot('true accuracy', (batch_true_correct / show_num))
vis.plot('cla_loss', fromDlpack(to_dlpack(batch_cla_loss / show_num)).item())
vis.plot('reg_loss', fromDlpack(to_dlpack(batch_reg_loss / show_num)).item())
vis.plot('box_loss', fromDlpack(to_dlpack(batch_box_loss / show_num)).item())
vis.plot('objective_loss', fromDlpack(to_dlpack(batch_objective_loss / show_num)).item())
vis.plot('aux_seg_loss', fromDlpack(to_dlpack(batch_aux_seg_loss / show_num)).item())
vis.plot('aux_offset_loss', fromDlpack(to_dlpack(batch_aux_offset_loss / show_num)).item())
vis.plot('accuracy', fromDlpack(to_dlpack(batch_correct / float(batch_num))).item())
vis.plot('true accuracy', fromDlpack(to_dlpack(batch_true_correct / show_num)).item())
batch_correct = 0.0
batch_cla_loss = 0.0
batch_reg_loss = 0.0
batch_box_loss = 0.0
batch_aux_seg_loss = 0.0
batch_aux_offset_loss = 0.0
batch_num = 0.0
batch_true_correct = 0.0
def tensor2cupy_gpu(torch_tensor: torch.Tensor) -> cupy.ndarray:
""" Convert a pytorch tensor to cupy array."""
with cupy.cuda.Device(torch_tensor.device.index):
cupy_array = fromDlpack(to_dlpack(torch_tensor))
return cupy_array