Combine two models to classify images

username: fedegonzal
cooked: <p>Hi, lets say I have two models:</p>
<ul>
<li>first one is trained to classify colors (red, green and blue)</li>
<li>the second one is trained to classify geometric shapes (rectangle, triangle and circle)</li>
</ul>
<p>I want to combine them to obtain a new model capable to understand a “red rectangle” and so one…</p>
<p>Which should be the proper way to do it?</p>
<p>Thanks!!</p>

username: KFrank
cooked: <p>Hi Federico!</p>
<aside class="quote no-group" data-username="fedegonzal" data-post="1" data-topic="164130" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fedegonzal/48/54235_2.png" class="avatar"> fedegonzal:</div>
<blockquote>
<p>I want to combine them to obtain a new model capable to understand a “red rectangle”</p>
</blockquote>
</aside>
<p>Plan A:</p>
<p>Let’s say that you have a <code>ColorClassifier</code> that does a (very) good job<br>
of classifying color, and a <code>ShapeClassifier</code> that does a (very) good job<br>
of classifying shape.</p>
<p>Just put your input to be classified into both separately.  If your<br>
<code>ColorClassifier</code> predicts “red” and your <code>ShapeClassifier</code> predicts<br>
“rectangle,” your “combined” classifier has predicted “red rectangle.”</p>
<p>If this works, just keep it simple and go with it.</p>
<p>Plan B:</p>
<p>Let’s say that the dataset you want to run your classifier on has some<br>
structure that links shape and color together somehow.  For example,<br>
perhaps rectangles are more likely to be red and triangles are more likely<br>
to be green, or perhaps you have images of red rectangles that also have<br>
green and blue blotches in the background.  You might then actually want<br>
to combine your classifiers into one.</p>
<p>Here’s the general approach I would take (making some plausible<br>
assumptions about the architectures of your classifiers):</p>
<p>Let’s assume that the final layer of <code>ColorClassifier</code> is a <code>Linear</code> with<br>
<code>out_features = 3</code>, and, for the sake of argument, with <code>in_features = 8</code>,<br>
that is, <code>Linear (8, 3)</code>, and, similarly, let’s assume that the final layer of<br>
<code>ShapeClassifier</code> is <code>Linear (12, 3)</code>.</p>
<p>So maybe the final section of <code>ColorClassifier</code> is:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.nn.Sequential (
   # various layers
   torch.nn.Linear (in_features = 128, out_features = 8),
   torch.nn.ReLU(),
   torch.nn.Linear (in_features = 8, out_features = 3)
)
</code></pre>
<p>and, similarly,  <code>ShapeClassifier</code> looks like:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.nn.Sequential (
   # various layers
   torch.nn.Linear (in_features = 256, out_features = 12),
   torch.nn.ReLU(),
   torch.nn.Linear (in_features = 12, out_features = 3)
)
</code></pre>
<p>That is, the second-to-last layer of <code>ColorClassifier</code> outputs a “feature<br>
vector” of length 8 (more precisely, a batch of feature vectors of shape<br>
<code>[nBatch, 8]</code>) and the second-to-last layer of <code>ShapeClassifier</code> outputs<br>
a feature vector of length 12.</p>
<p>The basic idea is that you would concatenate these two feature vectors<br>
together into a single feature vector of length 20, and then pass it through<br>
a final classification layer, <code>colorShapeLayer</code>.</p>
<p>The output of <code>colorShapeLayer</code> could be six values, the first three for<br>
the color classification, and the remaining three for the shape classification.<br>
But my intuition is that you would be better off having <code>colorShapeLayer</code><br>
output <em>nine</em> values, one each for each of the nine combinations of three<br>
colors and three shapes.</p>
<p>So your combined <code>ColorShapeClassifier</code> might look something like this:</p>
<pre data-code-wrap="python"><code class="lang-python">class ColorShapeClassifier (torch.nn.Module):
    def __init__ (self):
        super().__init__()
        self.colorFront = ColorFrontEnd()
        self.shapeFront = ShapeFrontEnd()
        self.colorShapeLayer = torch.nn.Linear (in_features = 20, out_features = 9)
    
    def forward (x):
        yColor = self.colorFront (x)
        yShape = self.shapeFront (y)
        yBoth = torch.cat ((yColor, yShape), dim = -1)
        z = torch.nn.functional.relu (yBoth)
        output = self.colorShapeLayer (z)
        return  output
</code></pre>
<p>where all but the last layer of <code>ColorClassifier</code> has been packaged as<br>
the <code>Module</code> <code>ColorFrontEnd</code>, and similarly for <code>ShapeFrontEnd</code>.</p>
<p><code>ColorShapeClassifier</code> can now be trained to make predictions for the<br>
nine classes that correspond to the nine combinations of the three colors<br>
and three shapes.</p>
<p>Initialize <code>ColorFrontEnd</code> and <code>ShapeFrontEnd</code> with pre-trained weights<br>
from <code>ColorClassifier</code> and <code>ShapeClassifier</code>, respectively (assumed<br>
to work well).  The idea is that the pre-trained <code>ColorFrontEnd</code> will output<br>
a feature vector that encodes useful information about color, and similarly<br>
for <code>ShapeFrontEnd</code>.  <code>colorShapeLayer</code> starts out initialized randomly.<br>
You would then <em>freeze</em> the pre-trained weights of <code>ColorFrontEnd</code> and<br>
<code>ShapeFrontEnd</code> and just train the weights of <code>colorShapeLayer</code> using<br>
your color-shape, nine-class dataset.</p>
<p>(After training <code>colorShapeLayer</code> for a while, you would probably want<br>
to fine-tune the entire <code>ColorShapeClassifier</code> model by unfreezing the<br>
“front-end” weights and training all of the weights jointly for a while.)</p>
<p>(It might also make sense to have more than one <code>Linear</code> layer after you<br>
<code>cat()</code> together the color and shape feature vectors or to use the output<br>
of the third-to-last layers of <code>ColorClassifier</code> and <code>ShapeClassifier</code> as<br>
your color and shape feature vectors that you <code>cat()</code> together.  I don’t know<br>
of any way to decide among such variations this other than experimenting<br>
with them.)</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: fedegonzal
cooked: <p>Thanks! I’ll try both, I’m using it to understand conceptual cases with simple practices</p>