PyTorch Conv2d使用复数类型报错：输入与偏置类型不匹配求解

官方文档说明Conv2d支持complex32、complex64、complex128复数类型，但在PyTorch 2.0.1+cu118环境中运行代码时，出现RuntimeError: Input type (c10::complex<float>) and bias type (float) should be the same报错，无法正常使用复数输入。

复现代码 #

import torch
import torch.nn as nn

print("Pytorch version")
print(torch.__version__)

m = nn.Conv2d(16, 33, 3, stride=2)

input = torch.randn(20, 16, 50, 100, dtype=torch.float32)
print(m(input))

print("NEXT")
input = torch.randn(20, 16, 50, 100, dtype=torch.complex64)
print(m(input))

运行报错信息 #

Pytorch version
2.0.1+cu118
tensor([[[[ 2.3293e-01, -3.9820e-01, -9.2724e-02,  ..., -3.7408e-02,
       -5.9544e-01,  1.1753e+00],
      ...,
      [-9.3071e-02, -8.7868e-02,  9.5579e-05,  ...,  7.3185e-01,
        1.3108e+00, -1.6092e-01]]]], grad_fn=<ConvolutionBackward0>)
NEXT
Traceback (most recent call last):
  File "/home/xxx/yyy/230914_learn_prop/test3.py", line 13, in <module>
print(m(input))
  File "/home/xxx/anaconda3/envs/piptorch20/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
  File "/home/xxx/anaconda3/envs/piptorch20/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 463, in forward
return self._conv_forward(input, self.weight, self.bias)
  File "/home/xxx/anaconda3/envs/piptorch20/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 459, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (c10::complex<float>) and bias type (float) should be the same

官方文档说明 #

This module supports complex data types i.e. complex32, complex64, complex128.

解决方案 #

• 转换参数为复数类型：初始化Conv2d后，手动将权重和bias转换为对应复数类型，示例代码：

  m = nn.Conv2d(16, 33, 3, stride=2)
  # 转换权重和bias为complex64
  m.weight = nn.Parameter(m.weight.to(torch.complex64))
  m.bias = nn.Parameter(m.bias.to(torch.complex64))
  

• 禁用bias：如果业务不需要bias，创建Conv2d时设置bias=False即可避免类型不匹配：

  m = nn.Conv2d(16, 33, 3, stride=2, bias=False)
  

内容的提问来源于stack exchange，提问作者modu01