mmegine如何实现在一个iter内分别计算两次loss并单独backward

[huaibovip]

如何优雅地使用mmengine实现下面代码的计算流程

for epoch in range(epoch_start, max_epoch):
    for data in train_loader:
        model.train()
        adjust_learning_rate(optimizer, epoch, max_epoch, lr)
        data = [t.cuda() for t in data]
        x, y = data[0], data[1]

        # step 1: x -> y
        x_in = torch.cat((x,y), dim=1)
        warped, flow = model(x_in)
        loss_ncc = criterion_ncc(warped, y) * weights[0]
        loss_reg = criterion_reg(flow) * weights[1]
        loss = loss_ncc + loss_reg
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # step 2: y -> x
        y_in = torch.cat((y, x), dim=1)
        warped, flow = model(y_in)
        def_seg = torch.cat(def_segs, dim=1)
        loss_ncc = criterion_ncc(warped, x) * weights[0]
        loss_reg = criterion_reg(flow) * weights[1]
        loss = loss_ncc + loss_reg
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

[huaibovip]

目前的解决方案：
继承TrainLoop，将 run_iter() 中

        outputs = self.runner.model.train_step(
            data_batch, optim_wrapper=self.runner.optim_wrapper)

修改为

        outputs = self.runner.model.train_step(
            data_batch, optim_wrapper=self.runner.optim_wrapper)

        # exchange_data() 用于交换 x 和 y
        data_batch = exchange_data(data_batch)
        outputs = self.runner.model.train_step(
            data_batch, optim_wrapper=self.runner.optim_wrapper)

目前，尚不知道这种做法存在有什么问题

[zhouzaida]

可以参考 https://mmengine.readthedocs.io/zh-cn/latest/examples/train_a_gan.html#id5

[huaibovip]

谢谢