Обучение автоэнкодеру на разных машинах

Question

Я пытаюсь настроить автоматический кодер для поддержки беспроводной передачи данных. Часть кодера будет жить на стороне передатчика приемопередатчика, а декодер будет жить на стороне приемника. Как правило, передатчик и приемник могут быть разделены на мили и будут существовать на разных компьютерах.

Автоэнкодер должен быть обучен с использованием реального физического канала, поэтому необходимо будет выполнять обратное распространение на двух разных компьютерах (компьютеры передатчика и приемника). У меня вопрос: как я могу запустить процесс обратного распространения на стороне приемника и завершить его на стороне передатчика?

Чтобы немного упростить эту проблему, если вы можете помочь мне выполнить backprop на двух разныхфайлы, которые могут быть достаточно для меня, чтобы расширить его по мере необходимости. Представьте, что кодировщик определяется одним файлом, а декодер - другим. Как мне выполнить backprop для этих двух отдельных файлов?

Я готов использовать либо pytorch, либо тензор потока, в зависимости от того, что лучше подходит для решения проблемы. Pytorch будет моим первым выбором, если это возможно.

Вот код pytorch для стандартного автоэнкодера, который находится в одном файле и действует на данные CIFAR. Вы можете увидеть, как обратное распространение выполняется в одной строке loss.backward (). Это то, что не будет работать, когда автокодер разделен между машинами.

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()