N
ETWORK -A GNOSTIC
K NOWLEDGE
T RANSFER
FOR
M EDICAL I MAGE S EGMENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Conventional transfer learning leverages weights of pre-trained networks, but mandates the need for similar neural architectures.
Alternatively, knowledge distillation can transfer knowledge between heterogeneous networks but often requires access to the original training data or additional generative networks.
Knowledge transfer between networks can be improved by being agnostic to the choice of network architecture and reducing the dependence on original training data.
We propose a knowledge transfer approach from a teacher to a student network wherein we train the student on an independent transferal dataset, whose annotations are generated by the teacher.
Experiments were conducted on ﬁve state-of-the-art networks for semantic segmentation and seven datasets across three imaging modalities.
We studied knowledge transfer from a single teacher, combination of knowledge transfer and ﬁne-tuning, and knowledge transfer from multiple teachers.
The student model with a single teacher achieved similar performance as the teacher; and the student model with multiple teachers achieved better performance than the teachers.
The salient features of our algorithm include: 1) no need for original training data or generative networks, 2) knowledge transfer between different architectures, 3) ease of implementation for downstream tasks by using the downstream task dataset as the transferal dataset, 4) knowledge transfer of an ensemble of models, trained independently, into one student model.
Extensive experiments demonstrate that the proposed algorithm is effective for knowledge transfer and easily tunable.
I NTRODUCTION
Deep learning often requires a sufﬁciently large training dataset, which is expensive to build and not easy to share between users.
For example, a big challenge with semantic segmentation of medical images is the limited availability of annotated data (Litjens et al., 2017).
Due to ethical concerns and conﬁdentiality constraints, medical datasets are not often released with the trained networks.
This highlights the need for knowledge transfer between neural networks, wherein the original training dataset does not need to be accessed.
On the other hand, according to the black-box metaphor in deep learning based methods, transferring knowledge is difﬁcult between heterogeneous neural networks.
To address these limitations, algorithms were proposed to reuse or share the knowledge of neural networks, such as network weight transfer (Tan et al., 2018), knowledge distillation (Hinton et al., 2015), federated learning (Yang et al., 2019), and self-training (Xie et al., 2020b).
Some conventional algorithms directly transfer the weights of standard large models that were trained on natural image datasets for different tasks (Kang & Gwak, 2019;
Motamed et al., 2019;
Jodeiri et al., 2019; Raghu et al., 2019).
For example, Iglovikov & Shvets (2018) adopted VGG pre-trained on ImageNet as the encoder of U-Net for 2D image segmentation.
Similarly, the convolutional 3D (Tran et al., 2015), pre-trained on natural video datasets, was used as the encoder of 3D U-Net for the 3D MR (Magnetic Resonance) medical image segmentation (Zeng et al., 2017).
Transferring the network weights generally requires adjustments to be made to the architecture of the receiver model, this in turn, limits the ﬂexibility of the receiver network.
Another technique that involves knowledge transfer is federated learning (Yang et al., 2019) ; it has received attention for its capability to train a large-scale model in a decentralized manner without requiring users’ data.
In general, federated learning approaches adopt the central model to capture the shared knowledge of all users by aggregating their gradients.
Due to the difﬁculties in transferring knowledge between heterogeneous networks, federated learning often requires all devices, including both the central servers and local users, to use the same neural architecture (Xie et al., 2020a).
To our best knowledge, there has been no federated learning system that uses heterogeneous networks.
Knowledge distillation is the process of transferring the knowledge of a large neural network or an ensemble of neural networks (teacher) to a smaller network (student) (Hinton et al., 2015).
Given a set of trained teacher models, one feeds training data to them and uses their predictions instead of the true labels to train the student model.
For effective transfer of knowledge, however, it is essential that a reasonable fraction of the training examples are observable by the student (Li et al., 2018) or the metadata at each layer is provided (Lopes et al., 2017).
Yoo et al.
(2019) used a generative network to extract the knowledge of a teacher network, which generated labeled artiﬁcial images to train another network.
As can be seen, Yoo et al.
’s method had to train an additional generative network for each teacher network.
Different from knowledge distillation, self-training aims to transfer knowledge to a more capable model.
The self-training framework (Scudder, 1965) has three main steps: train a teacher model on labeled images; use the teacher to generate pseudo labels on unlabeled images; and train a student model on the combination of labeled images and pseudo labeled images.
Xie et al.
(2020b) proposed self-training with a noisy student for classiﬁcation, which iterates this process a few times by treating the student as a teacher to relabel the unlabeled data and training a new student.
These studies required labeled images for training the student; moreover, they implicitly required that the pseudolabeled images should be similar in content to the original labeled images.
Inspired by self-training, we propose a network-agnostic knowledge transfer algorithm for medical image segmentation.
This algorithm transfers the knowledge of a teacher model to a student model by training the student on a transferal dataset whose annotations are generated by the teacher.
The algorithm has the following characteristics:
the transferal dataset requires no manual annotation and is independent of the teacher-training dataset; the student does not need to inherit the weights of the teacher, and such, the knowledge transfer can be conducted between heterogeneous neural architectures; it is straightforward to implement the algorithm with ﬁne-tuning to solve downstream task, especially by using the downstream task dataset as the transferal dataset; the algorithm is able to transfer the knowledge of an ensemble of models that are trained independently into one model.
We conducted extensive experiments on semantic segmentation using ﬁve state-of-the-art neural networks and seven datasets.
The neural networks include DeepLabv3+
(Chen et al., 2018), UNet (Ronneberger et al., 2015), AttU-Net (Oktay et al., 2018), SDU-Net (Wang et al., 2020), and
Panoptic-FPN (Kirillov et al., 2019).
Out of the seven datasets, four public datasets involve breast lesion, nerve structure, skin lesion, and natural image objects, and three internal/in-house datasets involve breast lesion (a single dataset with two splits) and thyroid nodule.
Experiments showed that the proposed algorithm performed well for knowledge transfer on semantic image segmentation.
A LGORITHM
The main goal of the proposed algorithm is to transfer the knowledge of one or more neural networks to an independent one, without access to the original training datasets.
This section presents the proposed knowledge transfer algorithm for semantic segmentation in Algorithm 1 and its application with ﬁne-tuning for a downstream task in Algorithm 2 (A.4).
The knowledge transfer algorithm ﬁrst employs one or more teacher models to generate pseudo masks for the transferal dataset and then trains the student model on the pseudo-annotated transferal dataset.
For an image x in the transferal dataset D , we get the pseudo mask by the weighted average of the teacher models’ outputs: y = (cid:80) T i ∈ T w i · T i ( x ) , where w i is the weight for the model T i .
The output of a model T i ( x ) is either a soft mask (pixel value varies from 0 to 1) or a binary mask (pixel value is 0 or 1).
Since the ensemble of teacher models is not our main focus, we simply set the weights equally for all teacher models.
We adopt two constraints to exclude images without salient targets, as shown in Figure 1.
The ﬁrst constraint is on the target pixel number in the pseudo mask, where target pixel indicates the pixel with a value above 0.5.
We exclude the image (of size 384 × 384 pixels) if the target pixel number is less than a threshold of 256.
The second constraint is on the gray-level entropy of the pseudo
Require:
Teacher set T with trained models, randomly initialized student model S
Require:
Transferal dataset D
Require: Target pixel number threshold α , gray-level entropy threshold β 1: for each datapoint x in D do 2:
Obtain pseudo mask y
=
(cid:80) T i ∈
T w i · T i ( x ) 3:
N ← number of target pixels (value above 0.5) in y 4:
E ← gray-level entropy of y 5: if N < α or E > β then 6:
Exclude x from D 7: else 8:
Update x as ( x, y ) 9: end if 10: end for 11: Train student model S on pseudo-annotated D 12: return Trained S mask.
A high entropy value implies the target and the background have no obvious difference and vice versa.
We exclude images with a gray-level entropy higher than a threshold of 2.5.
We employ the second constraint only in the case that the teacher model outputs a soft mask rather than a binary mask.
In this algorithm, the teacher models and student model can be independent of each other and do not rely on each other during the training or testing phase.
The transferal dataset is independent from the teacher-training datasets, which are not observable for the student model.
The student model gains the knowledge of the teacher by the transferal dataset.
As a result, the student model is trained to work well on datasets similar to the teacher-training dataset, while it is not aimed to predict the ground truth of the transferal dataset.
In the case that a small dataset D t with ground truth is available for a downstream segmentation task (target task), it is straightforward to further ﬁne-tune the student model that has been trained on pseudo-annotated transferal dataset, as Algorithm 2 shows.
The transferal dataset D is independent of the target task, while it can also be a non-annotated dataset from the target task.
R ELATED N EURAL N ETWORKS
We posit that if the teacher model is ineffective, a more capable student model can easily achieve similar performance; if the student model, however, is ineffective, the cause cannot be easily identiﬁed and the reason can be attributed to the knowledge transfer algorithm and the inherent limitations of the student model.
We experimented with ﬁve neural networks, all of which have different architectures and have been shown to provide the best outcome on various segmentation tasks.
DeepLabv3+: DeepLabv3+ (Chen et al., 2018) is one of Google’s latest and best performing semantic segmentation models.
It combines the advantages of spatial pyramid pooling with an encoderdecoder structure.
Spatial pyramid pooling captures rich contextual information while the decoder path is able to gradually recover object boundaries.
U-Net: U-Net (Ronneberger et al., 2015) has been the most widely used network for medical image segmentation.
U-Net consists of an encoder-decoder structure with the encoder path extracting rich semantic information, and the decoder path recovering resolution and enabling contextual information.
AttU-Net: AttU-Net (Oktay et al., 2018) is a modiﬁcation of U-Net, where attention gates are used to control the skip connections from the encoder to the decoder.
Attention gates make the model focus on target structures by highlighting important features and suppressing irrelevant features.
SDU-Net: SDU-Net (Wang et al., 2020) utilizes parallel atrous convolutions with different dilation rates in each layer, which is effective in capturing features in both small and large receptive ﬁelds.
SDU-Net has demonstrated better performance using merely 40 percent of the parameters in U-Net.
Panoptic-FPN: Panoptic-FPN (Kirillov et al., 2019) merges semantic segmentation and object detection, such that each pixel is given a class as in semantic segmentation and each object is given a unique ID as in object detection.
Panoptic-FPN provides a more rich and complete segmentation.
D ATASETS
Table 1 presents an overview of the seven image datasets that were used for the experiments.
One example for each dataset is presented in Figure 2.
XXX Breast Lesion: Patients evaluated and referred for Ultrasound (US)-guided biopsy (BI-RADS 4 or 5 mass) at XXX Institute between September 2017 and October 2019 were recruited.
Eligible subjects were adult patients with breast masses detected by conventional US imaging.
A total of 1385 US images were adopted and manually annotated by two radiologists.
These images were randomly assigned to XXX Breast Lesion-1 (n=1237) and XXX Breast Lesion-2 (n=148).
[Table]
Baheya Breast Lesion: The Baheya Breast Lesion dataset was released by Al-Dhabyani et al.
and it contains 780 breast US images among women between 25 and 75 years old.
Radiologists from Baheya Hospital, Egypt manually annotated the lesion boundaries.
The images are categorized into three classes, which are normal (n=133), benign (n=487), and malignant (n=210).
In our experiment, we only used images with benign or malignant lesions, totaling 697 images.
Thyroid Nodule: A total of 3,830 US images of thyroid nodules were retrospectively collected form the XXX Thyroid Nodule Dataset.
Data from patients (more than 18 years old) who underwent US examination before their thyroid biopsy between April 2010 and April 2012 were included.
Manual delineation of thyroid nodules by two expert radiologists were the reference standard.
Nerve Structure: Hosted on Kaggle, the US Nerve segmentation dataset formed the basis of a competition in segmenting nerve structures in US images of the neck.
The dataset includes potential noise, artifacts, and mistakes when labeling the ground truth as well as near-identical images due to the inherent image acquisition technique.
In our experiment, we utilized all 11,143 original images.
Skin Lesion: The skin-lesion dataset was from a challenge hosted by the International Skin Imaging
Collaboration (ISIC) in 2019 where the task was to perform multiclass classiﬁcation to distinguish dermoscopic images among nine distinct diagnostic categories of skin cancers.
We employed the original images of the training set with 25,331 images across 8 classes.
PASCAL VOC2012: Pascal VOC2012 challenge dataset is popular for object detection and segmentation tasks.
This dataset consists of natural images belonging to 20 different types of objects.
We adopted a total of 17,125 original images both for object detection and segmentation.
E XPERIMENT
Experimental evaluation was conducted in the space of semantic image segmentation.
The experiments included: knowledge transfer with a single teacher, combination of knowledge transfer and ﬁne-tuning, and knowledge transfer with multiple teachers.
As mentioned in Algorithms 1 and 2, we excluded images without salient targets in their pseudo masks.
So we also tested and compared the algorithm’s performance with and without image exclusion.
In addition, we evaluated the knowledge transfer capability of each transferal dataset, which is deﬁned as the ratio of the highest Dice score of all student models that are trained on the transferal dataset to the Dice score of the teacher model that generates pseudo masks for the transferal dataset.
It is obvious the student and the teacher Dice scores should be calculated on the same test dataset.
5. K NOWLEDGE T RANSFER WITH A S INGLE T EACHER
We adopted DeepLabv3+ as the teacher model with XXX Breast Lesion-1 as the teacher-training dataset, and employed U-Net, AttU-Net, SDU-Net, and Panoptic-FPN as the student models.
We used Baheya Breast Lesion, Thyroid Nodule, Nerve Structure, Skin Lesion, and PASCAL VOC as transferal datasets.
The trained teacher processed the transferal datasets to generate pseudo masks, which were used to train the students.
We tested the teacher and students on XXX Breast Lesionand Baheya Breast Lesion.
The performance was evaluated by computing the Dice score between the model output and manual annotations, i.e., ground truth.
[Table]
The teacher model had an average Dice score of 87.04% ( ± 16.01%) on XXX Breast Lesion-2 and an average Dice score of 65.93%( ± 32.55%) on Baheya Breast Lesion.
As shown in Tables 2 and 3, all student model obtained impressive knowledge.
For XXX Breast Lesion-2, SDU-Net and PanopticFPN obtained Dice scores similar to and even better than the teacher in most of the cases, even though the transferal datasets have different kinds of objects and even different image modalities, such as PASCAL VOC2012.
For Baheya Breast Lesion, all students trained on pseudo-annotated
Baheya Breast Lesion resulted in higher Dice scores than the teacher model even though the size of
Baheya Breast Lesion is much smaller than the teacher-training dataset.
The superior performance of the students could be partly due to domain adaptation, which is one of the advantages of our algorithm when the tansferral dataset is from the target task.
The performance of student models was improved by excluding images that had few target pixels or high gray-level entropy from the transferal dataset, especially when the transferal dataset modality was different from the teacher-training dataset.
It is because the image of a different modality may result in more pseudo masks with few target pixels or high gray-level entropy, which could be seen from the entropy distribution in Figure 3 (A.2).
However, image exclusion did not work that well on
Baheya Breast Lesion.
It is probably because Baheya Breast Lesion is similar to the teacher-training dataset and of a much smaller size; image exclusion could further harm its representation capability rather than improve its representation quality.
All the transferal datasets have strong knowledge transfer capabilities that are close to and even above 100%, except the Skin Lesion.
Although Skin Lesion did not perform as well as the other transferal datasets, its knowledge transfer capability achieved 81.16% and 87.30% according to the two tests.
In general it is not a challenge to build a transferal dataset with a strong knowledge transfer capability.
At the same time, it is necessary to point out that the four student networks have different learning abilities on the same transferal dataset.
However, it is common for different networks to learn differently even on the dataset with manual annotation.
5. C OMBINATION WITH F INE -T UNING
To evaluate the usefulness of knowledge transfer for downstream tasks, we conducted ﬁne-tuning using 50 images with ground truth annotations (from Baheya Breast Lesion and Thyroid Nodule, respectively).
Both the teacher model and student models that were trained in the previous experiments were further ﬁne-tuned.
For comparison, we also conducted the experiment by training the student models from scratch on these images with ground truth.
All the models were tested on the rest of the images of Baheya Breast Lesion and Thyroid Nodule, respectively
The ﬁne-tuned teacher model, DeepLabv3+, resulted in an average Dice score of 69.90% ± 31.01% on Baheya Breast Lesion and an average Dice score of 66.56% ± 27.78% on Thyroid Nodule.
The results of the student models are presented in Tables 4 and 5. All ﬁne-tuned student models outperformed the models trained from scratch by a large margin. With the same ﬁne-tuning dataset, some of the student models, such as SDU-Net and Panoptic-FPN, performed similar to or even better than
[Table] [Table]
20 ± 24.
37.15 ± 24.
46.88 ± 28.
52.88 ± 30.
Baheya Breast Lesion 59.29 ± 29.
65.31 ± 29.
70.03 ± 31.
70.61 ± 29.
Baheya Breast Lesion* 61.45 ± 29.
64.07 ± 30.
70.85 ± 28.
70.97 ± 30.
Thyroid Nodule 57.08 ± 28.
60.09 ± 30.
70.48 ± 30.
68.04 ± 29.
Thyroid Nodule* 54.06 ± 29.
59.38 ± 29.
71.60 ± 28.
68.60 ± 29.
56.58 ± 26.
63.21 ± 30.
66.55 ± 31.
71.19 ± 28.
Nerve Structure* 57.70 ± 28.
56.76 ± 29.
69.33 ± 29.
68.13 ± 30.
Skin Lesion 45.03 ± 26.
43.29 ± 27.
53.45 ± 29.
62.18 ± 31.
Skin Lesion* 52.41 ± 27.
51.71 ± 28.
70.62 ± 29.
67.73 ± 29.
PASCAL VOC 60.87 ± 30.
62.41 ± 30.
69.93 ± 29.
70.91 ± 28.
PASCAL VOC2012* 57.11 ± 29.
61.63 ± 30.
70.83 ± 29.
71.67 ± 29.
the ﬁne-tuned teacher model.
This further veriﬁed that the effect of knowledge transfer between heterogeneous neural networks and demonstrate that the proposed algorithm can be well applied to downstream task.
5. K NOWLEDGE T RANSFER WITH M ULTIPLE T EACHERS
We also tested if it was possible to transfer the knowledge of multiple weak independent teachers to build a stronger student.
We adopted U-Net, AttU-Net, SDU-Net, and Panoptic-FPN as the teacher models and employed DeepLabv3+ as a student model.
We train the teacher models on
XXX Breast Lesion-1 for a few epochs to ensure under-ﬁtting, thereby resulting in weak teachers.
Baheya Breast Lesion was used as the transferal dataset.
The trained teachers processed each image in the transferal dataset, with the average soft mask as pseudo mask.
Then we trained the student model on the pseudo-annotated transferal dataset.
We compared the performance of the teacher models and the student model on XXX Breast Lesion2 and Baheya Breast Lesion (Table 6).
The student model outperformed all the teacher models on both test datasets.
This is primarily due to fact that the student model actually learned from the ensemble of the teacher models, rather than any individual weak teacher model.
Since the process of generating ensemble networks is not the main focus of this study, we simply average their outputs.
With an advanced ensemble, we may be able to train the student model better.
[Table]
[Table] [Table]
D ISCUSSION
The salient features of our algorithm include: no need for original training data or extra generative networks; facilitating knowledge transfer between networks with different architectures; ease of implementation with ﬁne-tuning to solve downstream task especially when using the downstream task dataset as the transferal dataset; capability of knowledge transfer from multiple models, trained independently, into one student model.
The algorithm can also perform domain adaptation without extra computation in the case that the downstream task dataset is used as the transferal dataset.
To the best of our knowledge, this is the ﬁrst demonstration of image segmentation using pseudo annotations only.
A transferal dataset with a large number of images of various contents has higher possibility to capture rich targets in the pseudo masks, but it may also include many images without salient targets.
Future work may include optimization of a single transferal dataset or a combination of multiple transferal datasets to build a better one.
Different models may have different abilities to learn from the pseudo-annotated transferal dataset.
Understanding the differences in student model learning on pseudo-annotated transferal dataset and manually annotated dataset can help generate conﬁdence in the proposed algorithm.
Further study on knowledge transfer between neural networks with different input data structures or different number of prediction classes is also warranted.
Another possible direction includes extending the ensemble of teacher models to produce pseudo masks.
Medical imaging applications require interpretable algorithms which generate conﬁdence in the clinical user.
Knowledge transfer employs pseudo annotation for training; these pseudo annotations have no physical meaning.
It is imperative to examine and quantify the interpretability of student model before deploying models clinically.
C ONCLUSION
Knowledge transfer can be achieved between neural networks having different architectures and using an independent transferal dataset.
Knowledge transfer from multiple networks improves the performance of the student network over the teacher networks.
A A PPENDIX A. D ATA P RE - PROCESSING AND A UGMENTATION
In addition to pre-processing, we employed ﬁve methods for data augmentation.
Pre-processing: All the images were resized to 384 × 384 and the color images in Skin Lesion and
PASCAL VOC2012 were converted into gray-scale images.
Random Cropping: A percentage of consecutive rows and columns were cropped from the original image.
The
percentage was a random value in the range [70%, 100%] and [80%, 100%] for segmentation and classiﬁcation, respectively.
Horizontal Flipping:
The image was reversed horizontally, that is, from left to right.
Random Rotation: The image was rotated around its center.
The rotation angle was a random value in the range [-45 ◦ , 45 ◦ ] and [-30 ◦ , 30 ◦ ] for segmentation and classiﬁcation, respectively.
Gamma Adjustment : Gamma correction was conducted on an image by the gamma value, which was randomly sampled in the range [0.7, 1.3].
Contrast
Adjustment:
The
image
contrast of an
image
was adjusted by
the contrast factor , which was a random
value in the range
A. D ATA D ISTRIBUTION AND S ELECTION
We used DeepLabv3+ that was trained on XXX Breast Lesion-1 to process the images in the ﬁve transferal datasets including Baheya Breast Lesion, Thyroid Nodule, Nerve Structure, Skin Lesion, and PASCAL VOC2012, and obtained the pseudo mask.
Figure 3 presents the gray-level entropy distribution of the pseudo mask for each transferal dataset.
By excluding images without obvious targets, where the target pixel number is less than 256 and the gray-level entropy is over 2.5, we could reduce the dataset to a smaller one.
Table 7 presents the comparison of image numbers with and without image exclusion.
[Table]
A. T RAINING D ETAILS
Two common loss functions were adopted with the same weight for the segmentation tasks, the Dice loss function and binary cross entropy loss (BCE).
The Dice loss is deﬁned as:
[Equation]
where ( h, w ) represents the pixel coordinate, p h,w ∈ { 0 , 1 } is the mask ground truth, p h,w =indicates the pixel belonging to the target, 0 ≤ ˆ p h,w ≤ 1 is the prediction probability for the pixel belonging to the target, (cid:15) is a small real number.
The binary cross entropy loss is deﬁned as:
[Equation]
[Equation]
) where ( h, w ) represents the pixel coordinate, p h,w ∈ { 0 , 1 } is the mask ground truth, p h,w =indicates the pixel belonging to the target, 0 ≤ ˆ p h,w ≤ 1 is the prediction probability for the pixel belonging to the target, N is the number of pixels.
Adam optimizer was used for training all the networks.
The learning rate was set to 0.0001, and all the other parameters were set to the default values in PyTorch.
The batch size was set to 4.
The epoch number was set to 500 for the training on the teacher-training dataset.
Inversely proportional to the image number, the epoch numbers were set to 500, 175, 60, and 30 to for the training on pseudoannotated Baheya Breast Lesion, Thyroid Nodule, Nerve Structure, and Skin Lesion.
Moreover, the training epoch number for ﬁne-tuning was set to 50.
A. A LGORITHM OF K NOWLEDGE T RANSFER C OMBINED WITH F INE - TUNING
Algorithm 2 combines knowledge transfer and ﬁne-tuning for a downstream task of semantic image segmentation.
In the algorithm, the transferal dataset is independent of the downstream task, but if it is from the downstream task, the algorithm can beneﬁt from domain adaptation without extra computation.
Require:
Teacher set T with trained models, randomly initialized student model S
Require: Transferal dataset D , target task dataset D t with annotation
Require: Target pixel number threshold α , gray-level entropy threshold β 1: for each datapoint x in D do 2:
Obtain pseudo mask y
=
(cid:80) T i ∈
T w i · T i ( x ) 3:
N ← number of target pixels (value above 0.5) in y 4:
E ← gray-level entropy of y 5: if N < α or E > β then 6:
Exclude x from D 7: else 8:
Update x as ( x, y ) 9: end if 10: end for 11: Train student model S on pseudo-annotated D 12: Fine-tune student model S on target task dataset D t 13: return Trained student model S