Complex weights

rglenn: <p>I am trying to convert a unet module to use complex weights. Is there a list of layers that support complex data types? Do all of the layers and activation functions support complex types? How do I find out more information on complex weights and torch.</p>
<p>Thank you</p>

KFrank: <p>Hi Rachel!</p>
<aside class="quote no-group" data-username="rglenn" data-post="1" data-topic="203214" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/a88e57/48.png" class="avatar"> rglenn:</div>
<blockquote>
<p>I am trying to convert a unet module to use complex weights.</p>
</blockquote>
</aside>
<p>Good luck!  This is likely to be, at best, annoying.  It’s probably doable, with some<br>
work-arounds.</p>
<aside class="quote no-group">
<blockquote>
<p>Is there a list of layers that support complex data types?</p>
</blockquote>
</aside>
<p>Not that I am aware of.</p>
<aside class="quote no-group">
<blockquote>
<p>Do all of the layers and activation functions support complex types?</p>
</blockquote>
</aside>
<p>No.  Sometimes the functionality in question doesn’t make sense with complex<br>
numbers (for example, <code>ReLU</code>) and sometimes things that do make sense haven’t<br>
been implemented yet.</p>
<p>Complex support in pytorch is still a work in progress.  There are also some<br>
nuances working with complex tensors.  These are not pytorch’s fault, but you<br>
need to understand them and how they interact with your use case.</p>
<aside class="quote no-group">
<blockquote>
<p>How do I find out more information on complex weights and torch.</p>
</blockquote>
</aside>
<p>In general, look at pytorch’s documentation, but, in your particular context, your best<br>
bet will be to go through your complex U-Net step by step, see what still works, and<br>
implement work-arounds as necessary for what doesn’t.</p>
<p>Some comments:</p>
<p>Complex numbers are not <em>ordered</em>.  That is, for two complex numbers <code>u</code> and <code>v</code>,<br>
it doesn’t make sense to say <code>u &lt; v</code>.  This means that things like <code>ReLU</code> <em>can’t</em> work<br>
in a natural way for complex numbers.</p>
<p>Note, for example, that for this reason <code>MaxPool2d</code> – often used in convolutional<br>
neural networks – doesn’t work with complex tensors.</p>
<p>Pytorch’s autograd and optimizers have built into them the assumption that you<br>
are training your (complex) network with a <em>real</em> loss.  The current version of pytorch<br>
seems to be good at warning you if you don’t follow this rule, but this restriction has<br>
the potential to lead to some unexpected results.</p>
<p>Here are some illustrations of pytorch’s complex support and limitations:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.3.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt;
&gt;&gt;&gt; r = torch.ones (1, 1, 4, 4)
&gt;&gt;&gt; c = torch.ones (1, 1, 4, 4, dtype = torch.complex64)
&gt;&gt;&gt;
&gt;&gt;&gt; r
tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]]]])
&gt;&gt;&gt; c
tensor([[[[1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
          [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
          [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
          [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; lin = torch.nn.Linear (4, 2, dtype = torch.complex64)
&gt;&gt;&gt; conv = torch.nn.Conv2d (1, 1, 3, dtype = torch.complex64)
&gt;&gt;&gt;
&gt;&gt;&gt; relu = torch.nn.ReLU()            # won't work with complex -- complex not ordered
&gt;&gt;&gt; maxp = torch.nn.MaxPool2d (2)     # won't work with complex -- complex not ordered
&gt;&gt;&gt; avgp = torch.nn.AvgPool2d (2)     # should work with complex, but doesn't
&gt;&gt;&gt;
&gt;&gt;&gt; lin (c)
tensor([[[[-0.3101+0.3275j,  0.0007+0.5325j],
          [-0.3101+0.3275j,  0.0007+0.5325j],
          [-0.3101+0.3275j,  0.0007+0.5325j],
          [-0.3101+0.3275j,  0.0007+0.5325j]]]], grad_fn=&lt;ViewBackward0&gt;)
&gt;&gt;&gt; conv (c)
tensor([[[[-0.0212-0.0528j, -0.0212-0.0528j],
          [-0.0212-0.0528j, -0.0212-0.0528j]]]], grad_fn=&lt;AddBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; relu (r)
tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]]]])
&gt;&gt;&gt; maxp (r)
tensor([[[[1., 1.],
          [1., 1.]]]])
&gt;&gt;&gt; avgp (r)
tensor([[[[1., 1.],
          [1., 1.]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; # relu (r)                        # would raise error
&gt;&gt;&gt; # maxp (r)                        # would raise error
&gt;&gt;&gt; # avgp (r)                        # would raise error
&gt;&gt;&gt;
&gt;&gt;&gt; # lin (c).sum().backward()        # would raise error -- "loss" not real
&gt;&gt;&gt; # conv (c).sum().backward()       # would raise error -- "loss" not real
&gt;&gt;&gt;
&gt;&gt;&gt; lin (c).sum().abs().backward()    # works -- "loss" is real
&gt;&gt;&gt; lin.weight.grad
tensor([[-1.3540+3.7639j, -1.3540+3.7639j, -1.3540+3.7639j, -1.3540+3.7639j],
        [-1.3540+3.7639j, -1.3540+3.7639j, -1.3540+3.7639j, -1.3540+3.7639j]])
&gt;&gt;&gt;
&gt;&gt;&gt; conv (c).sum().abs().backward()   # works -- "loss" is real
&gt;&gt;&gt; conv.weight.grad
tensor([[[[-1.4866-3.7135j, -1.4866-3.7135j, -1.4866-3.7135j],
          [-1.4866-3.7135j, -1.4866-3.7135j, -1.4866-3.7135j],
          [-1.4866-3.7135j, -1.4866-3.7135j, -1.4866-3.7135j]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; # implement AvgPool2d with Conv2d
&gt;&gt;&gt;
&gt;&gt;&gt; avgp_conv = torch.nn.Conv2d (1, 1, 2, 2, bias = False, dtype = torch.complex64)
&gt;&gt;&gt; avgp_conv.weight.requires_grad = False
&gt;&gt;&gt; avgp_conv.weight.copy_ (torch.ones (1, 1, 2, 2, dtype = torch.complex64) / 4)
Parameter containing:
tensor([[[[0.2500+0.j, 0.2500+0.j],
          [0.2500+0.j, 0.2500+0.j]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; s = torch.randn (1, 1, 4, 4)
&gt;&gt;&gt; t = torch.randn (1, 1, 4, 4, dtype = torch.complex64)
&gt;&gt;&gt;
&gt;&gt;&gt; s
tensor([[[[ 0.7203, -1.4108, -0.4384,  0.3551],
          [ 0.3730, -1.3050, -0.7983,  1.0442],
          [-0.1227,  0.4022, -1.4295, -0.5656],
          [ 0.6971,  0.1258, -0.0434,  0.5366]]]])
&gt;&gt;&gt; t
tensor([[[[ 0.2242-1.3234j, -0.0633+0.5293j,  0.4339-0.3124j,  0.0696-0.1074j],
          [-0.3859+0.6665j,  0.2333+1.0793j, -0.1318+0.6501j, -0.2297+0.9898j],
          [-0.8779-0.0999j, -0.4100-0.2612j, -0.8163+0.8257j,  0.1332+0.9607j],
          [ 0.1305+0.8744j,  0.6930+0.1200j, -0.7032+0.7892j,  0.6485+1.2661j]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; avgp (s)
tensor([[[[-0.4056,  0.0407],
          [ 0.2756, -0.3755]]]])
&gt;&gt;&gt; avgp_conv (s.type (torch.complex64))
tensor([[[[-0.4056+0.j,  0.0407+0.j],
          [ 0.2756+0.j, -0.3755+0.j]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; avgp_conv (t)
tensor([[[[ 0.0021+0.2379j,  0.0355+0.3050j],
          [-0.1161+0.1583j, -0.1844+0.9604j]]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

rglenn: <p>Thank you, this really helps to clear up some things.</p>