SOLVED …
So @lesscomfortable was thinking on the right track when asking about batch size as hooks only store the last batch worth of data in hook.stored.
Here’s a couple of ways by which you could implement the hook described above and discussed in much more detail here.
Option 1
class StoreHook(HookCallback):
def on_train_begin(self, **kwargs):
super().on_train_begin(**kwargs)
self.outputs = []
def hook(self, m, i, o):
return o
def on_batch_end(self, train, **kwargs):
if (not train): self.outputs.append(self.hooks.stored[0])
data = ImageDataBunch.from_folder(PATH, ds_tfms=get_transforms(), size=img_sz, bs=bsz).normalize(imagenet_stats)
learn = create_cnn(data, models.resnet34, metrics=[error_rate])
# the last 2 layers are Dropout and Linear (for predictions)
nn_module = learn.model[-1][-3] # BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
learn.callbacks += [ StoreHook(learn, modules=flatten_model(nn_module)) ]
learn.fit_one_cycle(1)
cb = learn.callbacks[0]
torch.cat(cb.outputs).shape, len(data.valid_ds)
# => torch.Size([455, 512]), 455
Option 2
class StoreHook2(Callback):
def __init__(self, module):
super().__init__()
self.custom_hook = hook_output(module)
self.outputs = []
def on_batch_end(self, train, **kwargs):
if (not train): self.outputs.append(self.custom_hook.stored)
data = ImageDataBunch.from_folder(PATH, ds_tfms=get_transforms(), size=img_sz, bs=bsz).normalize(imagenet_stats)
learn = create_cnn(data, models.resnet34, metrics=[error_rate])
# the last 2 layers are Dropout and Linear (for predictions)
nn_module = learn.model[-1][-3] # BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
learn.callbacks += [ StoreHook2(nn_module) ]
learn.fit_one_cycle(1)
cb = learn.callbacks[0]
torch.cat(cb.outputs).shape, len(data.valid_ds)
# => torch.Size([455, 512]), 455
``