def lbp_validate(corpus, model, helper, device):
model.eval()
K, M = args.K, helper.markov_order
total_loss, nexamples = 0.0, 0
for i in range(len(corpus)):
batch = corpus[i].to(device)
T, bsz = batch.size()
if T <= 1:
continue
if T not in cache:
edges, nodeidxs, ne = get_hmm_stuff(T, M, K)
cache[T] = (edges, nodeidxs, ne)
edges, nodeidxs, ne = cache[T]
nedges = edges.size(0)
edges, ne = edges.to(device), ne.view(1, -1).to(device)
un_lpots = model.get_obs_lps(batch) # bsz x T x K log unary potentials
ed_lpots = model.get_edge_scores(edges,
T) # nedges x K*K log potentials
exed_lpots = ed_lpots.view(nedges, 1, K, K)
# get approximate unclamped marginals
nodebs, facbs, _, _, _ = dolbp(exed_lpots,
edges,
niter=args.lbp_iter,
renorm=True,
randomize=args.randomize_lbp,
tol=args.lbp_tol)
xnodebs, xfacbs, _, _, _ = dolbp(exed_lpots.expand(nedges, bsz, K, K),
edges,
x=batch,
emlps=un_lpots.transpose(0, 1),
niter=args.lbp_iter,
renorm=True,
randomize=args.randomize_lbp,
tol=args.lbp_tol)
# reshape log unary marginals: T x bsz x K -> bsz x T x K
tau_u = torch.stack([nodebs[t] for t in range(T)]).transpose(0, 1)
taux_u = torch.stack([xnodebs[t] for t in range(T)]).transpose(0, 1)
# reshape log fac marginals: nedges x bsz x K x K -> bsz x nedges x K x K
tau_e = torch.stack([facbs[e] for e in range(nedges)]).transpose(0, 1)
taux_e = torch.stack([xfacbs[e]
for e in range(nedges)]).transpose(0, 1)
# exponentiate
tau_u, tau_e = (tau_u.exp() + EPS).view(1,
-1), (tau_e.exp() + EPS).view(
1, -1)
taux_u, taux_e = (taux_u.exp() + EPS).view(
bsz, -1), (taux_e.exp() + EPS).view(bsz, -1)
fx, _, _, _ = bethe_fex(taux_u, taux_e, un_lpots.view(bsz, -1),
ed_lpots.view(1, -1).expand(bsz, -1),
ne.expand(bsz, -1))
fz, _, _, _ = bethe_fez(tau_u, tau_e, ed_lpots.view(1, -1), ne)
loss = fx - fz * bsz
total_loss += loss.item()
nexamples += bsz
return total_loss, nexamples
def validate_unsup_am(corpus, model, infnet, cache, penfunc, neginf, device,
args):
model.eval()
infnet.eval()
K, M = args.K, args.markov_order
total_out_loss, total_pen_loss, nexamples = 0.0, 0.0, 0
for i in range(len(corpus)):
batch = corpus[i].to(device)
T, bsz = batch.size()
if T <= 1: # annoying
continue
if T not in cache:
edges, nodeidxs, ne = get_hmm_stuff(T, M, K)
cache[T] = (edges, nodeidxs, ne)
edges, nodeidxs, ne = cache[T]
edges = edges.to(device) # symbolic edge representation
ne, nodeidxs = ne.view(1, -1).to(device), nodeidxs.to(
device) # 1 x T*K, # T x maxne
npenterms = (nodeidxs != 2 * edges.size(0)).sum().float()
# maximize wrt rho
ed_lpots = model.get_edge_scores(edges,
T) # nedges x K*K log potentials
pred_rho = infnet.q(edges, T)
un_lpots = model.get_obs_lps(batch) # bsz x T x K log unary potentials
pred_rho_x = infnet.qx(batch, edges, T)
out_loss, open_loss = outer_loss(pred_rho_x,
pred_rho,
un_lpots.view(bsz, -1),
ed_lpots.view(1, -1),
nodeidxs,
K,
ne,
neginf,
penfunc=penfunc)
total_out_loss += out_loss.item()
total_pen_loss += 1 / npenterms * open_loss.item()
nexamples += bsz
return total_out_loss, total_pen_loss, nexamples
args = parser.parse_args()
device = torch.device("cuda" if args.cuda else "cpu")
torch.manual_seed(args.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
ntypes, max_verts, max_seqlen = 10002, 3, 30
model = HybEdgeModel(ntypes, max_verts, args).to(device)
if args.infarch == "rnnnode":
infnet = RNodeInfNet(ntypes, max_seqlen, args).to(device)
else:
infnet = TNodeInfNet(ntypes, max_seqlen, args).to(device)
T, M, K = 10, 3, 30
edges, nodeidxs, ne = get_hmm_stuff(T, M, K)
edges, ne = edges.to(device), ne.view(1, -1).to(device)
nodeidxs = nodeidxs.to(device)
npenterms = (nodeidxs != 2 * edges.size(0)).sum().float()
nedges = edges.size(0)
EPS = 1e-38
neginf = torch.Tensor(1, 1, 1).fill_(-1e18).to(device)
if args.penfunc == "l2":
penfunc = lambda x, y: ((x - y) * (x - y)).sum(-1)
elif args.penfunc == "l1":
penfunc = lambda x, y: (x - y).abs().sum(-1)
elif args.penfunc == "js":
penfunc = lambda x, y: 0.5 * (batch_kl(x, y) + batch_kl(y, x))
elif args.penfunc == "kl1":
penfunc = lambda x, y: batch_kl(x, y)
def main(args, helper, cache, max_seqlen, max_verts, ntypes, trbatches,
valbatches):
print("main args", args)
torch.manual_seed(args.seed)
if torch.cuda.is_available():
if not args.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
device = torch.device("cuda" if args.cuda else "cpu")
if args.infarch == "rnnnode":
infctor = RNodeInfNet
else:
infctor = TNodeInfNet
model = HybEdgeModel(ntypes, max_verts, args).to(device)
if "exact" not in args.loss:
infnet = infctor(ntypes, max_seqlen, args).to(device)
bestmodel = HybEdgeModel(ntypes, max_verts, args)
if "exact" not in args.loss:
bestinfnet = infctor(ntypes, max_seqlen, args)
else:
bestinfnet = None
if args.penfunc == "l2":
penfunc = lambda x, y: ((x - y) * (x - y)).sum(-1)
elif args.penfunc == "l1":
penfunc = lambda x, y: (x - y).abs().sum(-1)
elif args.penfunc == "js":
penfunc = lambda x, y: 0.5 * (batch_kl(x, y) + batch_kl(y, x))
elif args.penfunc == "kl1":
penfunc = lambda x, y: batch_kl(x, y)
elif args.penfunc == "kl2":
penfunc = lambda x, y: batch_kl(y, x)
else:
penfunc = None
neginf = torch.Tensor(1, 1, 1).fill_(-1e18).to(device)
best_loss, prev_loss = float("inf"), float("inf")
lrdecay, pendecay = False, False
if "exact" in args.loss:
if args.optalg == "sgd":
popt1 = torch.optim.SGD(model.parameters(), lr=args.lr)
else:
popt1 = torch.optim.Adam(model.parameters(), lr=args.lr)
else:
if args.optalg == "sgd":
popt1 = torch.optim.SGD(model.parameters(), lr=args.lr)
popt2 = torch.optim.SGD(infnet.parameters(), lr=args.ilr)
else:
popt1 = torch.optim.Adam(model.parameters(), lr=args.lr)
popt2 = torch.optim.Adam(infnet.parameters(), lr=args.ilr)
if args.check_corr:
from utils import corr
# pick a graph to check
T, K = 10, args.K
edges, nodeidxs, ne = get_hmm_stuff(T, args.markov_order, K)
edges, ne = edges.to(device), ne.view(1, -1).to(device)
nodeidxs = nodeidxs.to(device)
npenterms = (nodeidxs != 2 * edges.size(0)).sum().float()
nedges = edges.size(0)
with torch.no_grad():
#un_lpots = model.get_obs_lps(batch) # bsz x T x K log unary potentials
ed_lpots = model.get_edge_scores(edges,
T) # nedges x K*K log potentials
exed_lpots = ed_lpots.view(nedges, 1, K, K)
# get approximate unclamped marginals
nodebs, facbs, _, _, _ = dolbp(exed_lpots,
edges,
niter=args.lbp_iter,
renorm=True,
randomize=args.randomize_lbp,
tol=args.lbp_tol)
tau_u = torch.stack([nodebs[t]
for t in range(T)]).transpose(0,
1) # 1 x T x K
tau_e = torch.stack([facbs[e] for e in range(nedges)
]).transpose(0, 1) # 1 x nedge x K x K
tau_u, tau_e = (tau_u.exp() + EPS), (tau_e.exp() + EPS)
for i in range(args.z_iter):
with torch.no_grad(
): # these functions are used in calc'ing the loss below too
pred_rho = infnet.q(edges, T) # nedges x K^2 logits
# should be 1 x T x K and 1 x nedges x K^2
predtau_u, predtau_e, _ = get_taus_and_pens(pred_rho,
nodeidxs,
K,
neginf,
penfunc=penfunc)
predtau_u, predtau_e = predtau_u.exp() + EPS, predtau_e.exp(
) + EPS
# i guess we'll just pick one entry from each
un_margs = tau_u[0][:, 0] # T
bin_margs = tau_e[0][:, K - 1, K - 1] # nedges
pred_un_margs = predtau_u[0][:, 0] # T
pred_bin_margs = predtau_e[0].view(nedges, K, K)[:, K - 1,
K - 1] # nedges
print(
i, "unary corr: %.4f, binary corr: %.4f" % (corr(
un_margs, pred_un_margs), corr(bin_margs, pred_bin_margs)))
popt2.zero_grad()
pred_rho = infnet.q(edges, T) # nedges x K^2 logits
in_loss, ipen_loss = inner_lossz(pred_rho.view(1, -1),
ed_lpots.view(1, -1), nodeidxs, K,
ne, neginf, penfunc)
in_loss = in_loss + args.pen_mult / npenterms * ipen_loss
print("in_loss", in_loss.item())
in_loss.backward()
clip_opt_params(popt2, args.clip)
popt2.step()
exit()
bad_epochs = -1
for ep in range(args.epochs):
if args.loss == "exact":
ll, ntokes = exact_train(trbatches, model, popt1, helper, device,
args)
print("Epoch {:3d} | train tru-ppl {:8.3f}".format(
ep, math.exp(-ll / ntokes)))
with torch.no_grad():
vll, vntokes = exact_validate(valbatches, model, helper,
device)
print("Epoch {:3d} | val tru-ppl {:8.3f}".format(
ep, math.exp(-vll / vntokes)))
# if ep == 4 and math.exp(-vll/vntokes) >= 280:
# break
voloss = -vll
elif args.loss == "lbp":
oloss, nex = lbp_train(trbatches, model, popt1, helper, device,
args)
print("Epoch {:3d} | train out_loss {:8.3f}".format(
ep, oloss / nex))
with torch.no_grad():
voloss, vnex = lbp_validate(valbatches, model, helper, device)
print("Epoch {:3d} | val out_loss {:8.3f} ".format(
ep, voloss / vnex))
else: # infnet
oloss, iloss, ploss, nex = train_unsup_am(trbatches, model, infnet,
popt1, popt2, cache,
penfunc, neginf, device,
args)
print("Epoch {:3d} | train out_loss {:.3f} | train in_loss {:.3f}".
format(ep, oloss / nex, iloss / nex))
with torch.no_grad():
voloss, vploss, vnex = validate_unsup_am(
valbatches, model, infnet, cache, penfunc, neginf, device,
args)
print(
"Epoch {:3d} | val out_loss {:.3f} | val barr_loss {:.3f}".
format(ep, voloss / vnex, vploss / vnex))
if args.loss != "exact":
with torch.no_grad():
# trull, ntokes = exact_validate(trbatches, model, helper, device)
# print("Epoch {:3d} | train tru-ppl {:.3f}".format(
# ep, math.exp(-trull/ntokes)))
vll, vntokes = exact_validate(valbatches, model, helper,
device)
print("Epoch {:3d} | val tru-ppl {:.3f}".format(
ep, math.exp(-vll / vntokes)))
voloss = -vll
# trppl = math.exp(-trull/ntokes)
# if (ep == 0 and trppl > 3000) or (ep > 0 and trppl > 1000):
# break
if voloss < best_loss:
best_loss = voloss
bad_epochs = -1
print("updating best model")
bestmodel.load_state_dict(model.state_dict())
if bestinfnet is not None:
bestinfnet.load_state_dict(infnet.state_dict())
if len(args.save) > 0 and not args.grid:
print("saving model to", args.save)
torch.save(
{
"opt":
args,
"mod_sd":
bestmodel.state_dict(),
"inf_sd":
bestinfnet.state_dict()
if bestinfnet is not None else None,
"bestloss":
bestloss
}, args.save)
if (voloss >= prev_loss or lrdecay) and args.optalg == "sgd":
for group in popt1.param_groups:
group['lr'] *= args.decay
for group in popt2.param_groups:
group['lr'] *= args.decay
#decay = True
if (voloss >= prev_loss or pendecay):
args.pen_mult *= args.pendecay
print("pen_mult now", args.pen_mult)
pendecay = True
prev_loss = voloss
if ep >= 2 and math.exp(best_loss / vntokes) > 650:
break
print("")
bad_epochs += 1
if bad_epochs >= 5:
break
if args.reset_adam: #bad_epochs == 1:
print("resetting adam...")
for group in popt2.param_groups:
group['lr'] *= args.decay # not really decay
# if args.reset_adam and ep == 1: #bad_epochs == 1:
# print("resetting adam...")
# popt2 = torch.optim.Adam(infnet.parameters(), lr=args.ilr)
return bestmodel, bestinfnet, best_loss
def lbp_train(corpus, model, popt, helper, device, args):
model.train()
K, M = args.K, helper.markov_order
total_loss, nexamples = 0.0, 0
niter, nxiter = 0, 0
perm = torch.randperm(len(corpus))
for i, idx in enumerate(perm):
popt.zero_grad()
batch = corpus[idx.item()].to(device)
T, bsz = batch.size()
# if T <= 1 or (M == 3 and T <= 3): # annoying
# continue
if T <= 1:
continue
if T not in cache:
edges, nodeidxs, ne = get_hmm_stuff(T, M, K)
cache[T] = (edges, nodeidxs, ne)
edges, nodeidxs, ne = cache[T]
nedges = edges.size(0)
edges, ne = edges.to(device), ne.view(1, -1).to(device)
un_lpots = model.get_obs_lps(batch) # bsz x T x K log unary potentials
ed_lpots = model.get_edge_scores(edges,
T) # nedges x K*K log potentials
with torch.no_grad():
exed_lpots = ed_lpots.view(nedges, 1, K, K)
# get approximate unclamped marginals
nodebs, facbs, ii, _, _ = dolbp(exed_lpots,
edges,
niter=args.lbp_iter,
renorm=True,
randomize=args.randomize_lbp,
tol=args.lbp_tol)
xnodebs, xfacbs, iix, _, _ = dolbp(exed_lpots.expand(
nedges, bsz, K, K),
edges,
x=batch,
emlps=un_lpots.transpose(0, 1),
niter=args.lbp_iter,
renorm=True,
randomize=args.randomize_lbp,
tol=args.lbp_tol)
niter += ii
nxiter += iix
# reshape log unary marginals: T x bsz x K -> bsz x T x K
tau_u = torch.stack([nodebs[t] for t in range(T)]).transpose(0, 1)
taux_u = torch.stack([xnodebs[t]
for t in range(T)]).transpose(0, 1)
# reshape log fac marginals: nedges x bsz x K x K -> bsz x nedges x K x K
tau_e = torch.stack([facbs[e]
for e in range(nedges)]).transpose(0, 1)
taux_e = torch.stack([xfacbs[e]
for e in range(nedges)]).transpose(0, 1)
# exponentiate
tau_u, tau_e = (tau_u.exp() + EPS).view(1, -1), (tau_e.exp() +
EPS).view(1, -1)
taux_u, taux_e = (taux_u.exp() + EPS).view(
bsz, -1), (taux_e.exp() + EPS).view(bsz, -1)
fx, _, _, _ = bethe_fex(taux_u, taux_e, un_lpots.view(bsz, -1),
ed_lpots.view(1, -1).expand(bsz, -1),
ne.expand(bsz, -1))
fz, _, _, _ = bethe_fez(tau_u, tau_e, ed_lpots.view(1, -1), ne)
loss = fx - fz * bsz
total_loss += loss.item()
loss.div(bsz).backward()
clip_opt_params(popt, args.clip)
popt.step()
nexamples += bsz
if (i + 1) % args.log_interval == 0:
print(
"{:5d}/{:5d} | its {:3.2f}/{:3.2f} | out_loss {:8.3f}".format(
i + 1, perm.size(0), niter / (i + 1), nxiter / (i + 1),
total_loss / nexamples))
return total_loss, nexamples
def train_unsup_am(corpus, model, infnet, moptim, ioptim, cache, penfunc,
neginf, device, args):
"""
opt is just gonna be for everyone
"""
model.train()
infnet.train()
K, M = args.K, args.markov_order
total_out_loss, total_in_loss, nexamples = 0.0, 0.0, 0
total_pen_loss = 0.0
perm = torch.randperm(len(corpus))
for i, idx in enumerate(perm):
batch = corpus[idx.item()].to(device)
T, bsz = batch.size()
if T <= 1: # annoying
continue
if T not in cache:
edges, nodeidxs, ne = get_hmm_stuff(T, M, K)
cache[T] = (edges, nodeidxs, ne)
edges, nodeidxs, ne = cache[T]
edges = edges.to(device) # symbolic edge representation
ne, nodeidxs = ne.view(1, -1).to(device), nodeidxs.to(
device) # 1 x T*K, # T x maxne
npenterms = (nodeidxs != 2 * edges.size(0)).sum().float()
# maximize wrt rho
with torch.no_grad():
ed_lpots = model.get_edge_scores(edges,
T) # nedges x K*K log potentials
# if args.reset_adam:
# ioptim = torch.optim.Adam(infnet.parameters(), lr=args.ilr)
for _ in range(args.z_iter):
ioptim.zero_grad()
pred_rho = infnet.q(edges, T) # nedges x K^2 logits
in_loss, ipen_loss = inner_lossz(pred_rho.view(1, -1),
ed_lpots.view(1, -1), nodeidxs, K,
ne, neginf, penfunc)
total_in_loss += in_loss.item() * bsz
total_pen_loss += args.pen_mult / npenterms * ipen_loss.item(
) * bsz
in_loss = in_loss + args.pen_mult / npenterms * ipen_loss
in_loss.backward()
clip_opt_params(ioptim, args.clip)
ioptim.step()
pred_rho = pred_rho.detach()
if args.loss == "alt3":
# min wrt rho_x
with torch.no_grad():
un_lpots = model.get_obs_lps(
batch) # bsz x T x K log unary potentials
for _ in range(args.zx_iter):
ioptim.zero_grad()
pred_rho_x = infnet.qx(batch, edges, T)
out_loss1, open_loss1 = inner_lossx(
pred_rho_x, un_lpots.view(bsz, -1),
ed_lpots.view(1, -1).expand(bsz, -1), nodeidxs, K,
ne.expand(bsz, -1), neginf, penfunc)
out_loss1 = out_loss1 + args.pen_mult / npenterms * open_loss1
total_pen_loss += args.pen_mult / npenterms * open_loss1.item()
out_loss1.div(bsz).backward()
clip_opt_params(ioptim, args.clip)
ioptim.step()
pred_rho_x = pred_rho_x.detach()
# min wrt params
moptim.zero_grad()
# even tho these don't change we needa do it again
un_lpots = model.get_obs_lps(batch) # bsz x T x K log unary potentials
ed_lpots = model.get_edge_scores(edges,
T) # nedges x K*K log potentials
if args.loss != "alt3": # jointly minimizing over rho_x
pred_rho_x = infnet.qx(batch, edges, T)
openfunc = penfunc
else:
openfunc = None
out_loss, open_loss = outer_loss(pred_rho_x,
pred_rho,
un_lpots.view(bsz, -1),
ed_lpots.view(1, -1),
nodeidxs,
K,
ne,
neginf,
penfunc=openfunc)
total_out_loss += out_loss.item()
if args.loss != "alt3":
total_pen_loss += args.pen_mult / npenterms * open_loss
out_loss = out_loss + args.pen_mult / npenterms * open_loss
out_loss.div(bsz).backward()
clip_opt_params(moptim, args.clip)
moptim.step()
nexamples += bsz
if (i + 1) % args.log_interval == 0:
print(
"{:5d}/{:5d} | out_loss {:8.5f} | in_loss {:8.5f} | pen_loss {:8.6f}"
.format(i + 1, perm.size(0), total_out_loss / nexamples,
total_in_loss / nexamples,
total_pen_loss / (nexamples * args.pen_mult)))
return total_out_loss, total_in_loss, total_pen_loss, nexamples