pytorch对自定义loss函数自动求梯度

最新推荐文章于 2024-06-25 10:43:41 发布

城俊BLOG

最新推荐文章于 2024-06-25 10:43:41 发布

阅读量1.3k

点赞数 1

分类专栏： Pytorch # DL-基础

本文链接：https://blog.csdn.net/qxqxqzzz/article/details/115877593

版权

DL-基础同时被 2 个专栏收录

83 篇文章 2 订阅

订阅专栏

Pytorch

47 篇文章 4 订阅

订阅专栏

通过 torch.autograd.grad

class MMD(nn.Module):
    def __init__(self):
        super(MMD, self).__init__()
        self.mmd = torch.nn.MSELoss()

    def forward(self,fc1Features1,fc1Features2):
        n = len(fc1Features1)
        fc1_1 = 1/n * torch.sum(fc1Features1,axis=0)
        fc1_2 = 1/n * torch.sum(fc1Features2,axis=0)
        fc1 = fc1_1 - fc1_2
        mmdLoss = torch.norm(fc1,p=2)
        mmdLoss = mmdLoss * mmdLoss
        return mmdLoss

def fc1_constrain(self, fc1Features1, fc1Features2, ):
    torch.cuda.current_stream().wait_stream(self.stream)
    mmdLoss = self.mmd.forward(fc1Features1,fc1Features2)
    # torch.autograd.grad（y,[x1,x2]） 返回y分别对x1和x2求得的偏导数
    grad1,grad2 = torch.autograd.grad(mmdLoss, [fc1Features1,fc1Features2], only_inputs=True)
    mmdGradList = [grad1,grad2]
    return mmdLoss, mmdGradList