跨不同机器训练自动编码器

Question

我正在尝试训练自动编码器以支持无线数据传输。编码器部分将位于收发器的发射器侧，解码器将位于接收器侧。通常，发射器和接收器可以相隔数英里，并且将存在于不同的计算机上。

自动编码器必须使用真实的物理通道进行训练，因此有必要在两台不同的计算机(发射器和接收器计算机(之间执行反向传播。我的问题是，如何在接收器侧开始反向传播过程，并在发射器侧完成？

为了使这个问题更简单一些，如果您可以帮助我在两个不同的文件中执行反向传播，这可能足以让我根据需要扩展它。假设编码器由一个文件定义，解码器由另一个文件定义。我将如何在这两个单独的文件中执行反向传播？

我愿意使用pytorch或tensorflow，以更适合解决问题的方式为准。如果可能的话，Pytorch将是我的首选。

这是标准自动编码器的pytorch代码，它位于一个文件中并作用于CIFAR数据。您可以看到反向传播是如何在一行丢失中执行的.back((。当自动编码器在机器之间拆分时，这是不起作用的。

import torch
import torchvision as tv
import torchvision.transforms as transforms
import torch.nn as nn
from torch.autograd import Variable
# Loading and Transforming data
transform = transforms.Compose([transforms.ToTensor(),  transforms.Normalize((0.4914, 0.4822, 0.4466), (0.247,            0.243, 0.261))])
trainTransform  = tv.transforms.Compose([tv.transforms.ToTensor(), tv.transforms.Normalize((0.4914, 0.4822, 0.4466), (0.247, 0.243, 0.261))])
trainset = tv.datasets.CIFAR10(root='./data',  train=True,download=True, transform=transform)
dataloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=False, num_workers=4)
testset = tv.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)
# Writing our model
class Autoencoder(nn.Module):
def __init__(self):
super(Autoencoder,self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(3, 6, kernel_size=5),
nn.ReLU(True),
nn.Conv2d(6,16,kernel_size=5),
nn.ReLU(True))
self.decoder = nn.Sequential(
nn.ConvTranspose2d(16,6,kernel_size=5),
nn.ReLU(True),
nn.ConvTranspose2d(6,3,kernel_size=5),
nn.ReLU(True),
nn.Sigmoid())
def forward(self,x):
x = self.encoder(x)
x = self.decoder(x)
return x
#defining some params
num_epochs = 5 #you can go for more epochs, I am using a mac
batch_size = 128
model = Autoencoder().cpu()
distance = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(),weight_decay=1e-5)
for epoch in range(num_epochs):
for data in dataloader:
img, _ = data
img = Variable(img).cpu()
# ===================forward=====================
output = model(img)
loss = distance(output, img)
# ===================backward====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
# ===================log========================
print('epoch [{}/{}], loss:{:.4f}'.format(epoch+1, num_epochs, loss.data.numpy()))

Answer 1

我已经简化了您的示例以更快地对其进行测试。这就是我想出的

import torch
import torch.nn as nn

class AE(nn.Module):
def __init__(self):
super(AE,self).__init__()
self.encoder = nn.Sequential(nn.Linear(20, 10),
nn.ReLU(True),
nn.Linear(10, 5),
nn.ReLU(True))
self.decoder = nn.Sequential(nn.Linear(5, 10),
nn.ReLU(True),
nn.Linear(10, 20),
nn.ReLU(True),
nn.Sigmoid())
torch.manual_seed(0)
batch_size = 2
input_size = 20
epochs = 3
model = AE()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters())
x = torch.randn(batch_size, input_size)
for i in range(epochs):
optimizer.zero_grad()
code = model.encoder(x)
# this would be where the code goes through your physical channel
torch.save(code, 'code.pt')
code_from_file = torch.load('code.pt')
# compute the decoding with the code recovered on the other side
reconstruction = model.decoder(code_from_file)
loss = criterion(reconstruction, x)
print("LOSS: ", loss)
# this runs the backward pass up to code_from_file
# because saving is non-differentiable and it has 
# no knowledge of the computational graph before `code_from_file`
loss.backward()  
# here you would move the gradient of code_from_file through the channel
torch.save(code_from_file.grad, 'code_grad.pt')
code_grad = torch.load('code_grad.pt')
# and recover it on the other side
# feed the gradient to `code.backward` that will run backward pass up to the input
code.backward(code_grad)
# now you have the gradients for the encoder part and you can step
optimizer.step()