Skip to content

Instantly share code, notes, and snippets.

@willwhitney
willwhitney / tree_stack.py
Last active December 6, 2023 01:54
utils for stacking and unstacking jax pytrees to deal with vmap
import numpy as np
from jax import numpy as jnp
from jax.lib import pytree
def tree_stack(trees):
"""Takes a list of trees and stacks every corresponding leaf.
For example, given two trees ((a, b), c) and ((a', b'), c'), returns
((stack(a, a'), stack(b, b')), stack(c, c')).
[{"x":0,"y":0},{"x":1,"y":1},{"x":2,"y":2},{"x":3,"y":3},{"x":4,"y":4},{"x":5,"y":5},{"x":6,"y":6},{"x":7,"y":7},{"x":8,"y":8},{"x":9,"y":9}]
class DmMujocoModel(nn.Module):
def __init__(self, embed_dim, env_name, traj_len, qpos_only=False, qpos_qvel=False):
super().__init__()
self.embed_dim = embed_dim
self.dataset = DmData(env_name, traj_len, qpos_only, qpos_qvel)
self.dataset.make_env()
self.env = self.dataset.env
self.dummy_parameter = nn.Parameter(torch.zeros(1))
def forward(self, s, a):
@willwhitney
willwhitney / handle_slurm_signals.py
Created August 24, 2018 19:56
Fragment of python code for catching signals from Slurm and restarting the job
import signal
# depends on requesting SIGUSR1 in runner file: https://gist.github.com/willwhitney/e1509c86522896c6930d2fe9ea49a522
def handle_signal(signal_value, _):
signame = signal.Signals(signal_value).name
if signal_value == signal.SIGUSR1:
print('Process {} got signal {}. Saving and restarting.'.format(
os.getpid(), signame), flush=True)
save_dynamics(epoch)
@willwhitney
willwhitney / load_tf.py
Created July 10, 2018 16:11
load tensorboard log files into pandas dataframes
from tensorboard.backend.event_processing import event_accumulator
import tensorflow as tf
import glob
import pandas as pd
tf.logging.set_verbosity(tf.logging.ERROR)
basedir = "/path/to/log/directory/"
def load_tf(dirname):
prefix = basedir + "tboard/VisibleSwimmer-v2/"
@willwhitney
willwhitney / runner.py
Last active September 27, 2020 09:48
Script for running grids of experiments on slurm
#----------------------------------------------
# things to change:
# code_dir (the full path of the directory that contains your source dir)
# true_source_dir (change it from TD3 to whatever your source dir is called)
# job_source_dir (someplace to throw a duplicate of the source dir for this job)
#----------------------------------------------
import os
import sys
import itertools
File "parallel_main.py", line 56, in <module>
parents = random.choices(population, scores + 1e-10, k=pop_size)
File "/home/wwhitney/anaconda3/lib/python3.6/multiprocessing/pool.py", line 260, in map
return self._map_async(func, iterable, mapstar, chunksize).get()
File "/home/wwhitney/anaconda3/lib/python3.6/multiprocessing/pool.py", line 608, in get
raise self._value
multiprocessing.pool.MaybeEncodingError: Error sending result:
<my object>
Reason: 'RuntimeError('unable to open shared memory object </torch_29253_1584213566> in read-write mode at /py/conda-bld/pytorch_1493680494901/work/torch/lib/TH/THAllocator.c:226',)'
Traceback (most recent call last):
File "/home/wwhitney/anaconda3/lib/python3.6/multiprocessing/process.py", line 249, in _bootstrap
self.run()
File "/home/wwhitney/anaconda3/lib/python3.6/multiprocessing/process.py", line 93, in run
self._target(*self._args, **self._kwargs)
File "/home/wwhitney/anaconda3/lib/python3.6/multiprocessing/pool.py", line 108, in worker
task = get()
File "/home/wwhitney/anaconda3/lib/python3.6/multiprocessing/queues.py", line 345, in get
return _ForkingPickler.loads(res)
File "/home/wwhitney/anaconda3/lib/python3.6/site-packages/torch/multiprocessing/reductions.py", line 86, in rebuild_storage_filename
@willwhitney
willwhitney / propublica_per_capita.js
Last active May 13, 2016 04:19
Getting per-state per capita usage numbers from https://projects.propublica.org/checkup with jank javascript.
/*
To use me, go to one of the individual drug pages from propublica (e.g. https://projects.propublica.org/checkup/drugs/1726),
open the console (press ⌘-⌥-J on Mac),
paste in this script and hit Enter.
A new column should appear in your table.
This is very much hacked together in the console, so... no guarantees it will work next week.
*/
states = [[1, "California", 39144818],
@willwhitney
willwhitney / himawari.py
Last active November 30, 2022 17:47 — forked from celoyd/hi8-fetch.py
Fetch and untile tiled Himawari-8 images from the http://himawari8.nict.go.jp PNG endpoint, then set them as desktop background on OSX
import requests
import sys
from datetime import datetime, timedelta
import pytz
from PIL import Image
from StringIO import StringIO
import os
import logging
# python himawari.py