两个模型共享部分参数,联合训练

想将model1的entities_emb,也当作第model2的entities_emb,这样在更新model2参数的时候,能更新entities_emb。

首先在定义模型时,模型二的参数中加入model1.entities_emb

model1 = Model1(n1, n2, norm=distance_norm, dim=embedding_dim).cuda()
model2 = Model(n3, n4, device, model1.entities_emb, args).to(
        device) 

然后各自训练:

for epoch in range(epochs):
     optimizer1.zero_grad()
     loss1 = criterion()
     loss1.backward(retain_graph=True)
     optimizer1.step()

     optimizer2.zero_grad()
     loss2 = criterion()
     loss2.backward(retain_graph=True)
     optimizer2.step()

思路就是这样了,成功实现联合训练

你可能感兴趣的:(小记,深度学习,人工智能)