CODEWITHSUNDEEP
tensorflowclassification
Blackberry
Blackberry

Reputation: 161

Changing a trained network to keep only a subset of its output

Suppose I have a trained TensorFlow classification network for 20 classes as in PASCAL VOC 2007: aeroplane, bicycle, ..., car, cat, ..., person, ..., tvmonitor.

Now, I would like to have a sub-network for only a subset of the classes, e.g., 3 classes: car, cat, person.

Then, I can use this network for testing or for re-training/fine-tuning on a new dataset, only for the 3 classes.

It should be possible to extract this sub-network out of the original network, since it is only the last layer that will change. We need to discard the neurons/weights for the discarded classes.

My question: Is there an easy way to do this in TensorFlow? It will be great if you can point to some sample code or similar solution. I have googled, but have not come across any mention of this.

The symmetric problem, expanding the number of classes without discarding the original weights, can potentially be useful for some people, but my current focus is the one above.

Upvotes: 0

Views: 431

Answers (2)

P-Gn
P-Gn

Reputation: 24581

If you want to only keep the output for a few slices, you could simply extract the corresponding slices from the last layer.

For example, let's assume the last layer is fully connected. Its weights are a tensor of size num_previous x num_output.

You want to keep only a few of these outputs, says output 1, 22, and 42. You can get the weights of your new fully connected layer as:

outputs_to_keep = [1, 22, 42]
new_W = tf.transpose(tf.gather(tf.transpose(old_W), outputs_to_keep))

Upvotes: 1

Manolo Santos
Manolo Santos

Reputation: 1913

It is possible to extract a pretrained subnet as you said. It is called transfer learning. There are different ways to do it, here you have one:

  1. Find the layer you want to start with. You can use Tensorboard to find it and then use graph.get_tensor_by_name() Usually you keep the convolutional layers and discard the fully connected ones.
  2. Connect your new layers (normally fully connected ones) to the previous layer.
  3. Freeze the variables (weights) of the pretrained layers using trainable=false. Alternatively, you can instruct the optimizer to update only the weights from the new layers.
  4. Train your model with the new classes.

Upvotes: 0

Related Questions