Return item in forward hook

I am working on multiple instance learning, and the first two steps before attempting to cluster images from a group are training a network (in this case a pre-trained DenseNet, of which I unfreeze the last few layers) on the images and then saving the encoded representation of these images. The issue is, despite catching in the right place, the forward hook I register does not actually return the value as I want it to.

Code:

def encoded_return(self, d_in, d_out):
    print(d_out.size())
    return d_out

def make_encoded_representations(studydir, network, verbose=False):
    # find all series in studydir
    DL = df.getDirectoryList(studydir)
    
    # Take all paths, and check if they contain series (they should)
    sis = 0
    for seriesdir in DL:
        
        # If it is an image series, files should be readable as a series
        sitkreader = sitk.ImageSeriesReader()

        # Check if there is a DICOM series in the dicom_directory
        series_IDs = sitkreader.GetGDCMSeriesIDs(seriesdir)
        if verbose:
            print ("Loading dicom folder %" + seriesdir)
            print ("Detected "+str(len(series_IDs))+" distinct series. Loading files ...")
        
        for idx, ID in enumerate(series_IDs):
            try:
                # Get all file names
                series_file_names = sitkreader.GetGDCMSeriesFileNames(seriesdir, series_IDs[idx])
                if verbose:
                    print(str(len(series_file_names))+" files in series. Attempting cleanup if necessary ...")
                file_sizes = []

                # Try cleaning out garbage from series
                for file in series_file_names:
                    filereader = sitk.ImageFileReader()
                    filereader.SetFileName(file)
                    tmp = filereader.Execute()
                    size = tmp.GetSize()
                    origin = tmp.GetOrigin()
                    spacing = tmp.GetSpacing()
                    file_sizes.append((size[0], size[1]))
                size_hist = Counter(file_sizes)
                wanted_size = max(size_hist, key=size_hist.get)
                series_file_names = [name for idx, name in enumerate(series_file_names) if file_sizes[idx] == wanted_size]
        
                # make representation
                tensor_representation = representation(series_file_names)
                # load network
                net = torch.load(network)
                # set to eval mode
                net.eval()
                # register forward hook to return an encoded representation and not the final classification result
                net.module.features.denseblock4.denselayer24.register_forward_hook(encoded_return)
                # let the network evaluate and grab encoded image
                encoded_image = net(tensor_representation)
                # save encoded image and class locally
                torch.save((encoded_image, int(-1)), seriesdir+'/'+str(int(sis)).zfill(3)+'_ER.pth')
                
                sis += 1
            except:
                if verbose:
                    print("Cannot make encoded representation for series "+str(ID)+" in dir "+str(seriesdir)+".")
                raise
        
    return None

Now, the print statement in the hook tells me there is a tensor of size(1,48,7,7) to be found, which is the one I want, yet the eventually saved tensor is of size(1,12), which is the original classification layer I added as part of the transfer learning in step 1. I thought that the return statement in the hook would terminate the entire forward pass of the network and yield me the desired tensor. How would I grab that tensor correctly?

No, the forward hook won’t terminate the execution and you could store the intermediate output in e.g. a list or dict as seen here.

1 Like

Ah, cheers! That does the trick.

(And I can finally delete the hacky garbage I put into place as a workaround hehe)