FIND_RFS = False # Optionally receptive fields, just do qualitative
GRAD_APPROX_N = None # Quality of gradient estimation (None for ideal)
ROTATION_FRAMES = 32 # Number of test images for gradients
# Limit on category set (speedup)
CAT_WHITELIST = None
GRAD_HEATMAPS = False # Plot RF heatmaps? (speedup)
BATCH_SIZE = 10 # Batch size for model runs
# Path parameters
DATA_OUT = Paths.data("exp/att_models/manual")
PLOT_OUT = Paths.plots("exp/att_models/manual")
ISO_PATH = Paths.data("imagenet/imagenet_iso224.h5")
FOUR_PATH = Paths.data('imagenet/imagenet_four224l0.h5')
# Instantiate the pyTorch model and select test units
model, ckpt = cornet.load_cornet("Z")
'''units = vx.random_voxels_for_model(
model, UNIT_LAYERS, N_UNIT, 3, 224, 224)'''
# Or reload an old set of units
'''rf_fname = os.path.join(DATA_OUT, "rfs_noatt.csv")
units, _ = lsq_fields.load_rf_csv(rf_fname)'''
units = lsq_fields.load_units_from_csv(os.path.join(DATA_OUT, UNIT_SET))
# ================================================= #
# Find RFs #
# ================================================= #
if FIND_RFS:
units = exp.gather_rfs(model, units, {}, BETAS, 224, ROTATION_FRAMES,
GRAD_APPROX_N, PLOT_OUT, DATA_OUT, GRAD_HEATMAPS)
def get_model(weight_root = Paths.data.join("models")):
from proc.cornet import load_cornet
return load_cornet('Z', model = CORnet_Z(), weight_root = weight_root)[0]
parser.add_argument("--channels",
type=int,
help='Number of channels in the input image. Default: 3')
parser.add_argument("--size",
type=int,
help='Size of the of input images (square). Default: 224')
parser.add_argument(
'--model_file',
default=None,
help='Optional python file with a function `get_model` returns a ' +
'pytorch model object.')
args = parser.parse_args()
args.layers = [eval('tuple(' + l + ')') for l in args.layers]
if args.model_file is None:
from proc import cornet
model, _ = cornet.load_cornet("Z")
else:
exec(open(args.model_file, 'r').read())
model = get_model()
units = vx.random_voxels_for_model(model, args.layers, args.n_unit,
args.channels, args.size, args.size)
unit_strs = vx.VoxelIndex.serialize(units)
with open(args.output_path, 'w') as f:
f.write('unit\n')
for l in unit_strs:
for s in unit_strs[l]:
f.write(s)
f.write('\n')
def main():
# Load input files
if args.model is None:
from proc import cornet
model, _ = cornet.load_cornet("Z")
else:
spec = importlib.util.spec_from_file_location("model", args.model)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)
model = model_module.get_model()
if args.regs is not None:
regs = det.load_logregs(args.regs)
if args.nodata:
four_task = det.FakeDetectionTask(args.cats, 224, 3)
else:
four_task = det.DistilledDetectionTask(args.comp_images_path,
whitelist=args.cats)
imgs, ys, _, img_ixs = four_task.val_set(None, args.test_n)
if args.attn is not None:
kws = atts.load_cfg(args.attn_cfg)
att_mods = atts.load_model(args.attn, **kws)
print("loading attention", args.attn)
else:
att_mods = {}
print("no attention", args.attn)
# Ensure model will be un on GPU if requested
if args.cuda:
model.cuda()
# Don't treat categories separately
# Setup outputs
fn_outputs = h5py.File(args.output_path, 'w')
# ====== Run behavior task ======
for c in four_task.cats:
print("Category:", c)
# Save true outputs so we don't have to load the dataset later
ys_ds = fn_outputs.create_dataset(f'{c}_y', (len(imgs[c]), ), np.int8)
ys_ds[...] = ys[c]
ys_ds.attrs['cat_names'] = np.array(list(regs.keys())).astype('S')
# record index of image in the dataset
ix_ds = fn_outputs.create_dataset(f'{c}_ix', (len(imgs[c]), ),
np.uint32)
ix_ds[...] = img_ixs[c]
# Create a place to store the eventual scores
fn_ds = fn_outputs.create_dataset(f'{c}_fn', (len(imgs[c]), ),
np.float32)
batches = (np.arange(len(imgs[c])), )
if args.batch_size > 0:
if len(imgs[c]) > args.batch_size:
n_batch = np.ceil(len(imgs[c]) / args.batch_size)
batches = np.array_split(batches[0], n_batch)
# NOTE: code assumes we'll still see images in expected order
for batch_n, batch_ix in enumerate(batches):
print(f"Batch {batch_n+1} / {len(batches)}")
run_batch(model, imgs[c], batch_ix, att_mods, fn_ds, regs[c])
cleanup()
fn_outputs.close()
def main():
# Load input files
if args.model is None:
from proc import cornet
model, _ = cornet.load_cornet("Z")
else:
spec = importlib.util.spec_from_file_location("model", args.model)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)
model = model_module.get_model()
units = lsq_fields.load_units_from_csv(args.unit_path)
if args.regs is not None:
regs = det.load_logregs(args.regs)
else:
regs = None
if args.nodata:
four_task = det.FakeDetectionTask(args.cats, 224, 3)
else:
four_task = det.DistilledDetectionTask(args.comp_images_path,
whitelist=args.cats)
imgs, ys, _ = four_task.val_set(None, args.test_n)
if args.attn is not None:
kws = atts.load_cfg(args.attn_cfg)
att_mods = atts.load_model(args.attn, **kws)
print("loading attention", args.attn)
else:
att_mods = {}
print("no attention", args.attn)
print(att_mods)
# Ensure model will be un on GPU if requested
if args.cuda:
model.cuda()
# Don't treat categories separately
cat_names = list(imgs.keys())
cat_ids = np.concatenate(
[np.repeat(i_cat, len(imgs[c])) for i_cat, c in enumerate(cat_names)])
cat_img_ixs = [ # Within-category image counter
np.cumsum(cat_ids == i) - 1 for i in range(len(cat_names))
]
cat_ns = {c: len(imgs[c]) for c in imgs}
ys = np.concatenate([a for a in ys.values()])
imgs = torch.cat([a for a in imgs.values()])
# Setup outputs
# We have to wait to initialize the hdf5 arrays until we
# know what size the wrt layers will be but can setup the
# file pointer now.
img_count = 0
if not args.no_back:
outputs = h5py.File(args.output_path, 'w')
else:
outputs = None
if args.regs is not None:
fn_output_fname = os.path.join(
os.path.dirname(args.output_path),
'fn_' + os.path.basename(args.output_path))
fn_outputs = h5py.File(fn_output_fname, 'w')
# Save image metadata along with behavioral outputs
ids_ds = fn_outputs.create_dataset('cat_ids', (cat_ids.size, ),
np.int8)
ids_ds[...] = cat_ids
ids_ds.attrs['cat_names'] = np.array(cat_names).astype('S')
ys_ds = fn_outputs.create_dataset('true_ys', (cat_ids.size, ), np.int8)
ys_ds[...] = ys
else:
fn_outputs = None
# Adjust batch size to processor if requested:
if args.batch_size == int(args.batch_size) or not args.cuda:
args.batch_size = int(args.batch_size)
else:
tot_mem = torch.cuda.get_device_properties(0).total_memory / (1024**3)
args.batch_size = int(args.batch_size * tot_mem)
print(f"[bp] tot_mem {tot_mem} batch_size {args.batch_size}")
# ====== Run backprop ======
batches = (np.arange(len(imgs)), )
if args.batch_size > 0:
if len(imgs) > args.batch_size:
n_batch = np.ceil(len(imgs) / args.batch_size)
batches = np.array_split(batches[0], n_batch)
# NOTE: code assumes we'll still see images in expected order
for batch_n, batch_ix in enumerate(batches):
if args.verbose:
print(f"Batch {batch_n+1} / {len(batches)}")
run_batch(model, imgs, batch_ix, units, att_mods, fn_outputs, outputs,
regs, cat_names, cat_ids, cat_img_ixs, cat_ns)
cleanup()
if not args.no_back:
outputs.close()
if args.regs is not None:
fn_outputs.close()
def setUp(self):
self.model, ckpt = cornet.load_cornet("Z")
self.task = det.IsolatedObjectDetectionTask(
Paths.data('imagenet.h5'), whitelist=['bathtub', 'artichoke'])
self.SIZE = 20
