Package: complexPyTorch
complexPyTorchA high-level toolbox for using complex valued neural networks in PyTorch.Before version 1.7 of PyTroch, complex tensor were not supported.
The initial version ofcomplexPyTorchrepresented complex tensor using two tensors, one for the real and one for the imaginary part.
Since version 1.7, compex tensors of typetorch.complex64are allowed, but only a limited number of operation are supported.
The current versioncomplexPyTorchuse complex tensors (hence requires PyTorch version >= 1.7) and add support for various operations and layers.InstallationpipinstallcomplexPyTorchComplex Valued Networks with PyTorchArtificial neural networks are mainly used for treating data encoded in real values, such as digitized images or sounds.
In such systems, using complex-valued tensor would be quite useless.
However, for physic related topics, in particular when dealing with wave propagation, using complex values is interesting as the physics typically has linear, hence more simple, behavior when considering complex fields.
complexPyTorch is a simple implementation of complex-valued functions and modules using the high-level API of PyTorch.
Following [C. Trabelsi et al., International Conference on Learning Representations, (2018)], it allows the following layers and functions to be used with complex values:LinearConv2dMaxPool2dRelu (ℂRelu)BatchNorm1d (Naive and Covariance approach)BatchNorm2d (Naive and Covariance approach)Citating the codeIf the code was helpful to your work, please consider citing it:Syntax and usageThe syntax is supposed to copy the one of the standard real functions and modules from PyTorch.
The names are the same as innn.modulesandnn.functionalexcept that they start withComplexfor Modules, e.g.ComplexRelu,ComplexMaxPool2dorcomplex_for functions, e.g.complex_relu,complex_max_pool2d.
The only usage difference is that the forward function takes two tensors, corresponding to real and imaginary parts, and returns two ones too.BatchNormFor all other layers, using the recommendation of [C. Trabelsi et al., International Conference on Learning Representations, (2018)], the calculation can be done in a straightforward manner using functions and modules formnn.modulesandnn.functional.
For instance, the functioncomplex_reluincomplexFunctions, or its associated moduleComplexReluincomplexLayers, simply performsreluboth on the real and imaginary part and returns the two tensors.
The complex BatchNorm proposed in [C. Trabelsi et al., International Conference on Learning Representations, (2018)] requires the calculation of the inverse square root of the covariance matrix.
This is implemented inComplexbatchNorm1DandComplexbatchNorm2Dbut using the high-level PyTorch API, which is quite slow.
The gain of using this approach, however, can be experimentally marginal compared to the naive approach which consists in simply performing the BatchNorm on both the real and imaginary part, which is available usingNaiveComplexbatchNorm1DorNaiveComplexbatchNorm2D.ExampleFor illustration, here is a small example of a complex model.
Note that in that example, complex values are not particularly useful, it just shows how one can handle complex ANNs.importtorchimporttorch.nnasnnimporttorch.nn.functionalasFfromtorchvisionimportdatasets,transformsfromcomplexPyTorch.complexLayersimportComplexBatchNorm2d,ComplexConv2d,ComplexLinearfromcomplexPyTorch.complexFunctionsimportcomplex_relu,complex_max_pool2dbatch_size=64trans=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,),(1.0,))])train_set=datasets.MNIST('../data',train=True,transform=trans,download=True)test_set=datasets.MNIST('../data',train=False,transform=trans,download=True)train_loader=torch.utils.data.DataLoader(train_set,batch_size=batch_size,shuffle=True)test_loader=torch.utils.data.DataLoader(test_set,batch_size=batch_size,shuffle=True)classComplexNet(nn.Module):def__init__(self):super(ComplexNet,self).__init__()self.conv1=ComplexConv2d(1,10,5,1)self.bn=ComplexBatchNorm2d(10)self.conv2=ComplexConv2d(10,20,5,1)self.fc1=ComplexLinear(4*4*20,500)self.fc2=ComplexLinear(500,10)defforward(self,x):x=self.conv1(x)x=complex_relu(x)x=complex_max_pool2d(x,2,2)x=self.bn(x)x=self.conv2(x)x=complex_relu(x)x=complex_max_pool2d(x,2,2)x=x.view(-1,4*4*20)x=self.fc1(x)x=complex_relu(x)x=self.fc2(x)x=x.abs()x=F.log_softmax(x,dim=1)returnxdevice=torch.device('cuda'iftorch.cuda.is_available()else'cpu')model=ComplexNet().to(device)optimizer=torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.9)deftrain(model,device,train_loader,optimizer,epoch):model.train()forbatch_idx,(data,target)inenumerate(train_loader):data,target=data.to(device).type(torch.complex64),target.to(device)optimizer.zero_grad()output=model(data)loss=F.nll_loss(output,target)loss.backward()optimizer.step()ifbatch_idx%100==0:print('Train Epoch:{:3}[{:6}/{:6}({:3.0f}%)]\tLoss:{:.6f}'.format(epoch,batch_idx*len(data),len(train_loader.dataset),100.*batch_idx/len(train_loader),loss.item()))# Run training on 50 epochsforepochinrange(50):train(model,device,train_loader,optimizer,epoch)AcknowledgmentsI want to thank Piotr Bialecki for his invaluable help on the PyTorch forum.