FCN с патчами создает границу

Question

Я пытаюсь обучить модель Unet делать предсказания регрессии на пиксель для изображений. Для этого я разделяю свое большое изображение (1000x1000) на квадраты 200x200 пикселей. Затем используйте это для обучения модели FCN с линейным конечным слоем. Функция потерь - потеря MSE. На этапе прогнозирования я извлекаю те же блоки, но склеиваю их вместе и получаю окончательное выходное изображение. Когда я это делаю, проблема, которую я получаю, заключается в том, что между границами ящиков есть разрывы. (Я четко вижу ящики)

Я пытался справиться с этим, подавая 250x250 ящиков в свой FCN и вычисляя потери для центральной области 200x200. Я делаю тот же процесс для состояния прогнозирования. Извлеките 250x250 фрагментов, обрежьте центральную область 200x200 и сожмите изображение обратно. Пожалуйста, посмотрите код ниже:

Функция потери:

criterion = nn.MSELoss()
optimizer = optim.Adam(self.model.parameters(), lr=LR)
for inputs, labels in train_loader:
    inputs, labels = inputs.to(device), labels.to(device)
    optimizer.zero_grad()
    output = model(inputs)
    output = output.squeeze()
     _, dimx, dimy = output.shape
     loss = criterion(output[:,25:dimx-25, 25:dimy-25], labels[:,25:dimx-25, 25:dimy-25])
     loss.backward()
     optimizer.step()

Мой код для прогнозов выглядит следующим образом:

pred = np.zeros((height, width))
for i in range(25, height, 200):
    for j in range(25, width, 200):
        patch = img[:, i-25:i+225, j-25:j+225]
        patch = torch.from_numpy(patch)
        patch = patch.unsqueeze(dim=0).to(device)
        out = model(patch)
        out = out[0,0,25:225, 25:225]
        pred[i:i+200, j:j+200] = out.cpu().numpy()

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

Answer 1

После некоторого устранения неполадок я понял, что у меня возникла эта проблема, потому что я выполнял пакетную нормализацию между каждым сверточным слоем. Удаление этого шага решило проблему разрыва.

Answer 2

Имеет смысл иметь разрыв рядом с границей, потому что не требуется, чтобы сеть имела плавные прогнозы для разных блоков во время обучения.

Я предполагаю, что у вас ограниченная память GPU, поэтому вы берете только 200x200 пикселейкак вход за один раз;Таким образом, я бы предложил следующие два возможных обходных пути:

Первый . Вы можете использовать torchvision.transform.RandomCrop , чтобы генерировать обрезанные области 200x200 в качестве входных данных для обучения. На этапе тестирования вы непосредственно вводите все изображение, чтобы сделать прогноз . Интуиция заключается в том, что модель может видеть полное разрешение изображений, которое совпадает с данными тестирования, при этом потребляя меньше памяти GPU во время обучения. В этом случае вы также можете ожидать, что модели требуется больше времени для изучения всех шаблонов обучающих данных, поскольку она видит только частичные данные за один раз. скажем, 0,5x, и сохранить размер вывода, то есть 1x. Например, в вашем случае, после понижающей дискретизации входного изображения до 200x200, модель использует его для прогнозирования меток уровня пикселей 1000x1000 (вы можете использовать билинейные слои с повышением частоты дискретизации или с понижением частоты). Этот обходной метод использовался в некоторых реализациях сегментации ( AdaptSeg , DISE ).