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): 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, size)) 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?