def train_stack(opt):
if opt["dataset"] == "omniglot":
stack_size = asl.util.repl(mnist_size, 0, opt["nchannels"])
ItemType = OmniGlot
dataloader = asl.util.omniglotloader
refresh_data = lambda dl: refresh_omniglot(dl, nocuda=opt["nocuda"])
else:
stack_size = asl.util.repl(mnist_size, 0, opt["nchannels"])
ItemType = Mnist
dataloader = asl.util.mnistloader
refresh_data = lambda dl: refresh_mnist(dl, nocuda=opt["nocuda"])
## Data structures and functions
class MatrixStack(asl.Type):
typesize = stack_size
tracegen = opt["tracegen"]
trace = tracegen(opt["nitems"], opt["nrounds"])
push = Push(MatrixStack, ItemType, arch=opt["arch"], arch_opt=opt["arch_opt"])
pop = Pop(MatrixStack, ItemType, arch=opt["arch"], arch_opt=opt["arch_opt"])
empty = ConstantNet(MatrixStack,
requires_grad=opt["learn_constants"],
init=opt["init"])
def ref_sketch(items, r, runstate):
return trace(items, r, runstate, push=list_push, pop=list_pop, empty=list_empty)
class StackSketch(asl.Sketch):
def sketch(self, items, r, runstate):
"""Example stack trace"""
return trace(items, r, runstate, push=push, pop=pop, empty=empty)
stack_sketch = StackSketch([List[OmniGlot]], [OmniGlot])
nstack = ModuleDict({"push": push,
"pop": pop,
"empty": empty,
"stack_sketch": stack_sketch})
# Cuda that shit
asl.cuda(nstack, opt["nocuda"])
# Hack to add random object as input to traces
def refresh_with_random(x):
# import pdb; pdb.set_trace()
refreshed = refresh_data(x)
return [refreshed[0], random.Random(0)]
# Loss
it = dataloader(opt["batch_size"], normalize=opt["normalize"])
loss_gen = asl.single_ref_loss(stack_sketch,
ref_sketch,
it,
refresh_with_random,
accum=opt["accum"])
if opt["learn_constants"]:
parameters = nstack.parameters()
else:
parameters = torch.nn.ModuleList([push, pop]).parameters()
return common.trainmodel(opt, nstack, loss_gen, parameters)
def train_set(opt):
trace = tracegen2(opt["nitems"], opt["nrounds"])
add = Add(MatrixSet, Mnist, arch=opt["arch"], arch_opt=opt["arch_opt"])
card = Card(MatrixSet, MatrixInteger, arch=opt["arch"], arch_opt=opt["arch_opt"])
empty = ConstantNet(MatrixSet,
requires_grad=opt["learn_constants"],
init=opt["init"])
def ref_sketch(items, seed, runstate):
return trace(items, seed, runstate, add=py_add, card=py_card, empty=py_empty_set)
class SetSketch(asl.Sketch):
def sketch(self, items, seed, runstate):
"""Example set trace"""
return trace(items, seed, runstate, add=add, card=card, empty=empty)
set_sketch = SetSketch([List[Mnist]], [Mnist])
nset = ModuleDict({"add": add,
"card": card,
"empty": empty,
"set_sketch": set_sketch})
# Cuda that shit
asl.cuda(nset, opt["nocuda"])
@dispatch(int)
def bridge(pyint):
width = mnist_size[1]
height = mnist_size[2]
res = asl.util.onehot2d(pyint, width, height, opt["batch_size"])
res = Variable(res, requires_grad=False)
return MatrixInteger(asl.cuda(res, opt["nocuda"]))
# Loss (horrible horrible hacking)
mnistiter = asl.util.mnistloader(opt["batch_size"])
# Make two random number objects which should stay in concert
model_random, ref_random = random.Random(0), random.Random(0)
model_inputs = [mnistiter, model_random]
ref_inputs = [mnistiter, ref_random]
def refresh_inputs(x):
mnist_dl, randomthing = x
a = asl.refresh_iter(mnist_dl, lambda x: Mnist(asl.util.image_data(x)))
return [a, random.Random(0)]
loss_gen = asl.single_ref_loss_diff_inp(set_sketch,
ref_sketch,
model_inputs,
ref_inputs,
refresh_inputs,
accum=opt["accum"])
if opt["learn_constants"]:
parameters = nset.parameters()
else:
parameters = torch.nn.ModuleList([add, card]).parameters()
return common.trainmodel(opt, nset, loss_gen, parameters)
def train_queue():
# Get options from command line
opt = asl.opt.handle_args(mnist_args)
opt = asl.opt.handle_hyper(opt, __file__)
class Enqueue(asl.Function, asl.Net):
def __init__(self="Enqueue", name="Enqueue", **kwargs):
asl.Function.__init__(self, [MatrixQueue, Mnist], [MatrixQueue])
asl.Net.__init__(self, name, **kwargs)
class Dequeue(asl.Function, asl.Net):
def __init__(self="Dequeue", name="Dequeue", **kwargs):
asl.Function.__init__(self, [MatrixQueue], [MatrixQueue, Mnist])
asl.Net.__init__(self, name, **kwargs)
nqueue = ModuleDict({
'enqueue':
Enqueue(arch=opt.arch, arch_opt=opt.arch_opt),
'dequeue':
Dequeue(arch=opt.arch, arch_opt=opt.arch_opt),
'empty':
ConstantNet(MatrixQueue)
})
queue_sketch = QueueSketch([List[Mnist]], [Mnist], nqueue, ref_queue())
asl.cuda(queue_sketch)
# Loss
mnistiter = asl.util.mnistloader(opt.batch_size)
loss_gen = asl.sketch.loss_gen_gen(
queue_sketch, mnistiter, lambda x: Mnist(asl.util.data.train_data(x)))
# Optimization details
optimizer = optim.Adam(nqueue.parameters(), lr=opt.lr)
asl.opt.save_opt(opt)
if opt.resume_path is not None and opt.resume_path != '':
asl.load_checkpoint(opt.resume_path, nqueue, optimizer)
asl.train(loss_gen,
optimizer,
maxiters=100000,
cont=asl.converged(1000),
callbacks=[
asl.print_loss(100), common.plot_empty, common.plot_observes,
asl.save_checkpoint(1000, nqueue)
],
log_dir=opt.log_dir)
def train_queue():
# Get options from command line
opt = asl.opt.handle_args(mnist_args)
opt = asl.opt.handle_hyper(opt, __file__)
class Enqueue(asl.Function, asl.Net):
def __init__(self="Enqueue", name="Enqueue", **kwargs):
asl.Function.__init__(self, [MatrixQueue, Mnist], [MatrixQueue])
asl.Net.__init__(self, name, **kwargs)
class Dequeue(asl.Function, asl.Net):
def __init__(self="Dequeue", name="Dequeue", **kwargs):
asl.Function.__init__(self, [MatrixQueue], [MatrixQueue, Mnist])
asl.Net.__init__(self, name, **kwargs)
nqueue = ModuleDict({'enqueue': Enqueue(arch=opt.arch,
arch_opt=opt.arch_opt),
'dequeue': Dequeue(arch=opt.arch,
arch_opt=opt.arch_opt),
'empty': ConstantNet(MatrixQueue)})
queue_sketch = QueueSketch([List[Mnist]], [Mnist], nqueue, ref_queue())
asl.cuda(queue_sketch)
# Loss
mnistiter = asl.util.mnistloader(opt.batch_size)
loss_gen = asl.sketch.loss_gen_gen(queue_sketch,
mnistiter,
lambda x: Mnist(asl.util.data.train_data(x)))
# Optimization details
optimizer = optim.Adam(nqueue.parameters(), lr=opt.lr)
asl.opt.save_opt(opt)
if opt.resume_path is not None and opt.resume_path != '':
asl.load_checkpoint(opt.resume_path, nqueue, optimizer)
asl.train(loss_gen, optimizer, maxiters=100000,
cont=asl.converged(1000),
callbacks=[asl.print_loss(100),
common.plot_empty,
common.plot_observes,
asl.save_checkpoint(1000, nqueue)],
log_dir=opt.log_dir)
def neural_queue(element_type, queue_type):
enqueue_img = EnqueueNet(queue_type, element_type)
dequeue_img = DequeueNet(queue_type, element_type)
empty_queue = ConstantNet(queue_type)
neural_ref = ModuleDict({
"enqueue": enqueue_img,
"dequeue": dequeue_img,
"empty": empty_queue
})
cuda(neural_ref)
return neural_ref
def benchmark_copy_sketch(batch_size, stack_len, seq_len, arch, log_dir, lr,
arch_opt, **kwargs):
stack_len = stack_len
seq_len = seq_len # From paper: between 1 and 20
BernSeq = bern_seq(seq_len)
MatrixStack = matrix_stack(1, seq_len, seq_len)
arch_opt['combine_inputs'] = lambda xs: stretch_cat(
xs, MatrixStack.size, 2)
arch_opt['activation'] = F.sigmoid
nstack = ModuleDict({
'push':
PushNet(MatrixStack, BernSeq, arch=ConvNet, arch_opt=arch_opt),
'pop':
PopNet(MatrixStack, BernSeq, arch=ConvNet, arch_opt=arch_opt),
'empty':
ConstantNet(MatrixStack)
})
refstack = ref_stack()
copy_sketch = CopySketch(BernSeq, nstack, refstack, seq_len)
cuda(copy_sketch)
bern_iter = BernSeq.iter(batch_size)
def loss_gen():
# Should copy the sequence, therefore the output should
items = take(bern_iter, seq_len)
rev_items = items.copy()
rev_items.reverse()
outputs = copy_sketch(items)
log("outputs", outputs)
log("items", items)
log("rev_items", rev_items)
# import pdb; pdb.set_trace()
return vec_dist(outputs, rev_items, dist=nn.BCELoss())
optimizer = optim.Adam(copy_sketch.parameters(), lr)
train(
loss_gen,
optimizer,
cont=converged(1000),
callbacks=[
print_loss(100),
every_n(plot_sketch, 500),
# common.plot_empty,
# common.plot_observes,
save_checkpoint(1000, copy_sketch)
],
log_dir=log_dir)
def train_clevrgen(opt):
class AddObject(asl.Function, asl.Net):
"Add object to scene"
def __init__(self, name="AddObject", **kwargs):
asl.Function.__init__(self, [Scene, Object], [Scene])
asl.Net.__init__(self, name, **kwargs)
class Render(asl.Function, asl.Net):
"Render a scene to an image"
def __init__(self, name="Render", **kwargs):
asl.Function.__init__(self, [Scene], [Image])
asl.Net.__init__(self, name, **kwargs)
class GenObject(asl.Function, asl.Net):
"Sample an object from noise"
def __init__(self, name="GenObject", **kwargs):
asl.Function.__init__(self, [Noise], [Object])
asl.Net.__init__(self, name, **kwargs)
nclevrgen = ModuleDict({'add_object': AddObject(arch=opt.arch,
arch_opt=opt.arch_opt),
'render': Render(arch=opt.arch,
arch_opt=opt.arch_opt),
'gen_object': Render(arch=opt.arch,
arch_opt=opt.arch_opt),
'empty_scene': ConstantNet(Scene)})
class Discriminate(asl.Function, asl.Net):
"Is ``Image`` from data distribution?"
def __init__(self, name="Discriminate", **kwargs):
asl.Function.__init__(self, [Image], [Probability])
asl.Net.__init__(self, name, **kwargs)
discriminate = Discriminate(arch=opt.arch, arch_opt=opt.arch_opt)
class ClevrGen(asl.Sketch):
"""Generate clevr image from noise"""
def sketch(self, noise):
"""Generate clevr image"""
# Add object 1
#FIXME this is a hack
noisetensor = asl.cuda(Variable(torch.rand((1,) + SMALL_IMG_SIZE)), opt.nocuda)
nnoise = Noise(expand_to_batch(noisetensor, opt.batch_size))
(object1, ) = nclevrgen.gen_object(nnoise)
scene = nclevrgen.empty_scene
(scene, ) = nclevrgen.add_object(scene, object1)
# Add object 2
(object2, ) = nclevrgen.gen_object(nnoise)
(scene, ) = nclevrgen.add_object(scene, object2)
# Add object 3
(object3, ) = nclevrgen.gen_object(nnoise)
(scene, ) = nclevrgen.add_object(scene, object3)
(img, ) = nclevrgen.render(scene)
asl.observe(img, 'rendered_img')
return (img, )
def ref_img_gen(img_iter):
img = next(img_iter)
asl.observe(img, 'rendered_img')
return (img, )
@dispatch(Image, Image)
def dist(x, y):
return discriminate(x.value, y.value)
clevrgen_sketch = ClevrGen([Noise], [Image])
# Cuda everything
asl.cuda(nclevrgen, opt.nocuda)
asl.cuda(clevrgen_sketch, opt.nocuda)
# Loss
img_dl = clevr_img_dl(opt.batch_size, normalize=False)
loss_gen = asl.ref_loss_gen(clevrgen_sketch,
ref_img_gen,
img_dl,
lambda ten: Image(asl.cuda(Variable(ten),
opt.nocuda)))
# Optimization details
parameters = nclevrgen.parameters()
optimizer = optim.Adam(parameters, lr=opt.lr)
asl.opt.save_opt(opt)
if opt.resume_path is not None and opt.resume_path != '':
asl.load_checkpoint(opt.resume_path, nclevrgen, optimizer)
asl.train(loss_gen, optimizer, maxiters=100000,
cont=asl.converged(1000),
callbacks=[asl.print_loss(100),
# common.plot_empty,
common.log_observes,
common.plot_observes,
# common.plot_internals,
asl.save_checkpoint(1000, nclevrgen)],
log_dir=opt.log_dir)
def train_clevrgen(opt):
class AddObject(asl.Function, asl.Net):
"Add object to scene"
def __init__(self, name="AddObject", **kwargs):
asl.Function.__init__(self, [Scene, Object], [Scene])
asl.Net.__init__(self, name, **kwargs)
class Render(asl.Function, asl.Net):
"Render a scene to an image"
def __init__(self, name="Render", **kwargs):
asl.Function.__init__(self, [Scene], [Image])
asl.Net.__init__(self, name, **kwargs)
class GenObject(asl.Function, asl.Net):
"Sample an object from noise"
def __init__(self, name="GenObject", **kwargs):
asl.Function.__init__(self, [Noise], [Object])
asl.Net.__init__(self, name, **kwargs)
nclevrgen = ModuleDict({
'add_object':
AddObject(arch=opt.arch, arch_opt=opt.arch_opt),
'render':
Render(arch=opt.arch, arch_opt=opt.arch_opt),
'gen_object':
Render(arch=opt.arch, arch_opt=opt.arch_opt),
'empty_scene':
ConstantNet(Scene)
})
class Discriminate(asl.Function, asl.Net):
"Is ``Image`` from data distribution?"
def __init__(self, name="Discriminate", **kwargs):
asl.Function.__init__(self, [Image], [Probability])
asl.Net.__init__(self, name, **kwargs)
discriminate = Discriminate(arch=opt.arch, arch_opt=opt.arch_opt)
class ClevrGen(asl.Sketch):
"""Generate clevr image from noise"""
def sketch(self, noise):
"""Generate clevr image"""
# Add object 1
#FIXME this is a hack
noisetensor = asl.cuda(
Variable(torch.rand((1, ) + SMALL_IMG_SIZE)), opt.nocuda)
nnoise = Noise(expand_to_batch(noisetensor, opt.batch_size))
(object1, ) = nclevrgen.gen_object(nnoise)
scene = nclevrgen.empty_scene
(scene, ) = nclevrgen.add_object(scene, object1)
# Add object 2
(object2, ) = nclevrgen.gen_object(nnoise)
(scene, ) = nclevrgen.add_object(scene, object2)
# Add object 3
(object3, ) = nclevrgen.gen_object(nnoise)
(scene, ) = nclevrgen.add_object(scene, object3)
(img, ) = nclevrgen.render(scene)
asl.observe(img, 'rendered_img')
return (img, )
def ref_img_gen(img_iter):
img = next(img_iter)
asl.observe(img, 'rendered_img')
return (img, )
@dispatch(Image, Image)
def dist(x, y):
return discriminate(x.value, y.value)
clevrgen_sketch = ClevrGen([Noise], [Image])
# Cuda everything
asl.cuda(nclevrgen, opt.nocuda)
asl.cuda(clevrgen_sketch, opt.nocuda)
# Loss
img_dl = clevr_img_dl(opt.batch_size, normalize=False)
loss_gen = asl.ref_loss_gen(
clevrgen_sketch, ref_img_gen, img_dl,
lambda ten: Image(asl.cuda(Variable(ten), opt.nocuda)))
# Optimization details
parameters = nclevrgen.parameters()
optimizer = optim.Adam(parameters, lr=opt.lr)
asl.opt.save_opt(opt)
if opt.resume_path is not None and opt.resume_path != '':
asl.load_checkpoint(opt.resume_path, nclevrgen, optimizer)
asl.train(
loss_gen,
optimizer,
maxiters=100000,
cont=asl.converged(1000),
callbacks=[
asl.print_loss(100),
# common.plot_empty,
common.log_observes,
common.plot_observes,
# common.plot_internals,
asl.save_checkpoint(1000, nclevrgen)
],
log_dir=opt.log_dir)
def train_stack(opt):
if opt["dataset"] == "omniglot":
stack_size = asl.util.repl(mnist_size, 0, opt["nchannels"])
ItemType = OmniGlot
dataloader = asl.util.omniglotloader
refresh_data = lambda dl: refresh_omniglot(dl, nocuda=opt["nocuda"])
else:
stack_size = asl.util.repl(mnist_size, 0, opt["nchannels"])
ItemType = Mnist
dataloader = asl.util.mnistloader
refresh_data = lambda dl: refresh_mnist(dl, nocuda=opt["nocuda"])
## Data structures and functions
class MatrixStack(asl.Type):
typesize = stack_size
tracegen = opt["tracegen"]
trace = tracegen(opt["nitems"], opt["nrounds"])
push = Push(MatrixStack,
ItemType,
arch=opt["arch"],
arch_opt=opt["arch_opt"])
pop = Pop(MatrixStack,
ItemType,
arch=opt["arch"],
arch_opt=opt["arch_opt"])
empty = ConstantNet(MatrixStack,
requires_grad=opt["learn_constants"],
init=opt["init"])
def ref_sketch(items, r, runstate):
return trace(items,
r,
runstate,
push=list_push,
pop=list_pop,
empty=list_empty)
class StackSketch(asl.Sketch):
def sketch(self, items, r, runstate):
"""Example stack trace"""
return trace(items, r, runstate, push=push, pop=pop, empty=empty)
stack_sketch = StackSketch([List[OmniGlot]], [OmniGlot])
nstack = ModuleDict({
"push": push,
"pop": pop,
"empty": empty,
"stack_sketch": stack_sketch
})
# Cuda that shit
asl.cuda(nstack, opt["nocuda"])
# Hack to add random object as input to traces
def refresh_with_random(x):
# import pdb; pdb.set_trace()
refreshed = refresh_data(x)
return [refreshed[0], random.Random(0)]
# Loss
it = dataloader(opt["batch_size"], normalize=opt["normalize"])
loss_gen = asl.single_ref_loss(stack_sketch,
ref_sketch,
it,
refresh_with_random,
accum=opt["accum"])
if opt["learn_constants"]:
parameters = nstack.parameters()
else:
parameters = torch.nn.ModuleList([push, pop]).parameters()
return common.trainmodel(opt, nstack, loss_gen, parameters)