Confused by output of "hook_output"

Part 1 (2019) Advanced (Part 1 v3)
4

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
``
9 Likes