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.


def encoded_return(self, d_in, d_out):
    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):
                # 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()
                    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
                # register forward hook to return an encoded representation and not the final classification result
                # let the network evaluate and grab encoded image
                encoded_image = net(tensor_representation)
                # save encoded image and class locally
      , int(-1)), seriesdir+'/'+str(int(sis)).zfill(3)+'_ER.pth')
                sis += 1
                if verbose:
                    print("Cannot make encoded representation for series "+str(ID)+" in dir "+str(seriesdir)+".")
    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)