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?