def build_nnet(dims, activation_fn=torch.nn.ReLU):
nnet = []
domains, codomains = parse_vnorms()
if args.learn_p:
if args.mixed:
domains = [torch.nn.Parameter(torch.tensor(0.)) for _ in domains]
else:
domains = [torch.nn.Parameter(torch.tensor(0.))] * len(domains)
codomains = domains[1:] + [domains[0]]
for i, (in_dim, out_dim, domain,
codomain) in enumerate(zip(dims[:-1], dims[1:], domains,
codomains)):
nnet.append(activation_fn())
nnet.append(
base_layers.get_linear(
in_dim,
out_dim,
coeff=args.coeff,
n_iterations=args.n_lipschitz_iters,
atol=args.atol,
rtol=args.rtol,
domain=domain,
codomain=codomain,
zero_init=(out_dim == 2),
))
return torch.nn.Sequential(*nnet)
def build_nnet(dims, activation_fn=torch.nn.ReLU):
nnet = []
domains, codomains = parse_vnorms()
for i, (in_dim, out_dim, domain,
codomain) in enumerate(zip(dims[:-1], dims[1:], domains,
codomains)):
if i > 0:
nnet.append(activation_fn())
nnet.append(
base_layers.get_linear(
in_dim,
out_dim,
coeff=args.coeff,
n_iterations=args.n_lipschitz_iters,
atol=args.sn_tol,
rtol=args.sn_tol,
domain=domain,
codomain=codomain,
zero_init=(out_dim == data_dim),
))
return torch.nn.Sequential(*nnet)
