Внедрение и запуск 3D U-сети для удаления шума из синтетических данных

Question

Я создаю сеть для шумоподавления некоторых синтетических изображений с течением времени. Мои данные (256,256,128,1), что составляет 256 * 256 изображений в градациях серого за 128 временных шагов (траектория - случайный косинус). У меня 120 изображений, поэтому весь набор данных имеет форму (120,256,256,128,1). Я добавил 5% случайного шума к этим данным для получения набора данных с шумом. Я использую TF2 и отношусь к нему больше как керас, так как я только начал.

Вот сводка моих данных:

x_train (120, 256, 256, 128, 1)
x_train_noisy (120, 256, 256, 128, 1)
x_validation (10, 256, 256, 128, 1)
x_validation_noisy (10, 256, 256, 128, 1)
x_test (10, 256, 256, 128, 1)
x_test_noisy (10, 256, 256, 128, 1)
x_train_input (100, 256, 256, 128, 1)
x_train_noisy_input (100, 256, 256, 128, 1)

Теперь я хочу создать трехмерную сеть для придания ей шума. Где я даю зашумленные изображения в качестве входных данных, а оригинал демонизируется как желаемый вывод и хочу, чтобы он изучал преобразование. Я новичок в этом, поэтому я использовал простой 2d unet из github в качестве шаблона и адаптировал его.

-Вот моя U-сеть: я хочу знать (1) Правильна ли эта архитектура U-сети? (что-то явно не так с этим и т. д.?) (это в основном нормальная сеть, как у Ронненберга).

inputs = Input((256,256,128,1))

conv1 = Conv3D(64, (3,3,3), activation='relu', padding='same')(inputs)
conv1 = Conv3D(64, (3,3,3), activation='relu', padding='same')(conv1)
pool1 = MaxPooling3D(pool_size=(2,2,2))(conv1)

conv2 = Conv3D(128, (3,3,3), activation='relu', padding='same')(pool1)
conv2 = Conv3D(128, (3,3,3), activation='relu', padding='same')(conv2)
pool2 = MaxPooling3D(pool_size=(2,2,2))(conv2)

conv3 = Conv3D(256, (3,3,3), activation='relu', padding='same')(pool2)
conv3 = Conv3D(256, (3,3,3), activation='relu', padding='same')(conv3)
pool3 = MaxPool3D(pool_size= (2,2,2))(conv3)

conv4 = Conv3D(512, (3,3,3), activation='relu', padding='same')(pool3)
conv4 = Conv3D(512, (3,3,3), activation='relu', padding='same')(conv4)

up4 = concatenate([Conv3DTranspose(256, (3,3,3), strides=(2,2,2), padding ='same')(conv4), conv3], axis=-1)

conv5 = Conv3D(256, (3,3,3), activation='relu', padding='same')(up4)
conv5 = Conv3D(256, (3,3,3), activation='relu', padding='same')(conv5)

up5 = concatenate([Conv3DTranspose(128, (3,3,3), strides=(2,2,2), padding='same')(conv5), conv2], axis=-1)

conv6 = Conv3D(128, (3,3,3), activation='relu', padding='same')(up5)
conv6 = Conv3D(128, (3,3,3), activation='relu', padding='same')(conv6)

up6 = concatenate([Conv3DTranspose(64, (3,3,3), strides=(2,2,2), padding ='same')(conv6), conv1], axis=-1)

conv7 = Conv3D(64, (3,3,3), activation='relu', padding='same')(up6)
conv7 = Conv3D(64, (3,3,3), activation='relu', padding='same')(conv7)

conv8 = Conv3D(1, (3,3,3), activation=None, padding='same')(conv7)

conv9 = Activation('relu')(conv8)

#now the keras Model and compilation
unet = Model(inputs = [inputs], outputs = [conv9])
unet.compile(optimizer = 'adam', loss = 'mean_squared_error')

-Теперь я пытаюсь соответствовать:

    save_dir = "/Users/Jack/Desktop/unet_50" 

es_cb = EarlyStopping(monitor='val_loss', patience=2, verbose=1, mode='auto')
chkpt = save_dir + 'unet_50_synthetic_denoise_weights.{epoch:02d}-{loss:.2f}-{val_loss:.2f}.hdf5'
cp_cb = ModelCheckpoint(filepath = chkpt, monitor='val_loss', verbose=1, save_best_only=True, mode='auto')
tensor_board = tf.keras.callbacks.TensorBoard(
                log_dir = "my_log_dir_unet_50",
                histogram_freq = 1,
                embeddings_freq = 1,)

history = unet.fit(x_train_noisy_input, x_train_input,
                   batch_size = 4,
                   epochs = 100,
                   verbose=1,
                   validation_data=(x_validation_noisy, x_validation),
                   callbacks = [es_cb, cp_cb],
                   shuffle=True)

-Но я получаю это, и оно вот так глохло около часа. Я почти уверен, что в эпоху должно быть что-то вроде индикатора прогресса:

Train on 100 samples, validate on 10 samples
Epoch 1/100

-Так вот мои вопросы: Что я делаю не так, что подгонка застопорилась (это должно занять несколько часов, а индикатор выполнения просто отстает). Я знаю, что на моем компьютере (ЦП) определенно работает что-то тяжелое.

-Вы видите что-то не так с кодом u-net?

- Что касается размера партии и разделения данных для проверки, видите ли вы очевидную ошибку или лучший способ выбора чисел?

Answer 1

Размер входных данных, которые вы подаете в свою сеть (256x256x128 изображений), огромен, кроме того, у вас есть 64 слоя на первом уровне вашей архитектуры.Я полагаю, что только с учетом извилистых слоев первого уровня все уже должно объединяться в нечто вроде 10–100 ГБ памяти GPU, которая слишком велика.

Я думаю, именно поэтому она застряла.удивлен тем, что он не передает никаких сообщений об ошибках, таких как переполнение памяти графического процессора.

Я бы порекомендовал уменьшить количество слоев первых уровней и изменить их на 2d (3,3,1) ядра (ядро будет работать только с размерами изображений, а не с измерением времени).