<jupyter_start><jupyter_text>Baseline  Initialize<jupyter_code>gpus = tf.config.list_physical_devices("GPU")
print(gpus)
if len(gpus) == 1:
    strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
else:
    strategy = tf.distribute.MirroredStrategy()
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": True})
print("Mixed precision enabled")<jupyter_output>Mixed precision enabled<jupyter_text>Data<jupyter_code>NUM_FOLDS = 5
SIZE = 120
BATCH_SIZE = 256
LR = 1e-3
EPOCHS = 10
INPUT_SHAPE = (SIZE, SIZE, 4)
PATH = Path("../input")
TRAIN_PATH = PATH / "train"
TEST_PATH = PATH / "test"
DAMS = (6071, 6304, 7026, 7629, 7767, 8944, 11107)
train_paths = list(TRAIN_PATH.glob("*.npy"))
filenames = [path.name for path in train_paths]
len(filenames)<jupyter_output><empty_output><jupyter_text>Visualize<jupyter_code>cmap = plt.cm.get_cmap("RdBu")
cmap = cmap.reversed()
def visualize(img, cmap=cmap):
    fig, axes = plt.subplots(1, 5, figsize=(20, 20))
    for i, ax in enumerate(axes[:-1]):
        ax.imshow(img[:, :, i], cmap=cmap)
    axes[-1].imshow(img[:, :, -1], cmap=cmap)
    plt.show()
idx = np.random.randint(len(train_paths))
img = np.load(train_paths[idx])
print(idx)
visualize(img)
img = np.load(train_paths[idx])
img = img[:, :, 0]
fig, ax = plt.subplots(figsize=(5, 5))

for dam in DAMS:
    x = dam // SIZE
    y = dam % SIZE
    circle = Circle((y, x), 1, color="r")
    ax.add_patch(circle)

ax.imshow(img)
plt.show()<jupyter_output><empty_output><jupyter_text>Make folds<jupyter_code>df = pd.DataFrame({"filename": filenames})
df = train_df.sample(frac=1).reset_index(drop=True)
df["fold"] = -1
kf = model_selection.KFold(n_splits=NUM_FOLDS)
for fold, (train_index, val_index) in enumerate(kf.split(df)):
    df.loc[val_index, "fold"] = fold
df.to_csv("../input/train_folds.csv", index=False)<jupyter_output><empty_output><jupyter_text>Dataset<jupyter_code>class DataGenerator(tf.keras.utils.Sequence):
    def __init__(self, df, batch_size, shuffle=False, test=False):
        self.filenames = df.filename.values
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.test = test
        self.on_epoch_end()

    def __len__(self):
        return math.ceil(len(self.filenames) / self.batch_size)

    def __getitem__(self, index):
        filenames = self.filenames[
            index * self.batch_size : (index + 1) * self.batch_size
        ]
        if self.test:
            x = self._data_generation(filenames)
            return x
        else:
            x, y = self._data_generation(filenames)
            return x, y

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.filenames)

    def _data_generation(self, filenames):
        x = np.zeros((len(filenames), SIZE, SIZE, 4))
        y = np.zeros((len(filenames), SIZE, SIZE, 1))
        for i, filename in enumerate(filenames):
            data = np.load((TEST_PATH if self.test else TRAIN_PATH) / filename)
            image = data[:, :, :4]
            #             image = image / 255
            #             image = image.astype(np.float32)
            x[i] = image

            if not self.test:
                target = data[:, :, -1].reshape(SIZE, SIZE, 1)
                target = np.where(target < 0, 0, target)
                #                 target = target / 255
                #                 target = target.astype(np.float32)
                y[i] = target

        if self.test:
            return x
        else:
            return x, y
# df = pd.read_csv(PATH / "train_folds.csv")
# gen = DataGenerator(df, batch_size=BATCH_SIZE, shuffle=True)
# for batch in gen:
#     x, y = batch
#     inp = x[10]
#     targ = y[10]
#     img = np.concatenate([inp, targ], axis=-1)
#     visualize(img, cmap=cmap)
#     break<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>df = pd.read_csv(PATH / "train_folds.csv")
def get_baseline(input_shape):
    inp = Input(input_shape)

    conv1 = Conv2D(64, 3, activation="relu", padding="same")(inp)
    bn1 = BatchNormalization()(conv1)
    pool1 = MaxPooling2D(2)(bn1)

    conv2 = Conv2D(128, 3, activation="relu", padding="same")(pool1)
    bn2 = BatchNormalization()(conv2)
    pool2 = MaxPooling2D(2)(bn2)

    convm = Conv2D(256, 3, activation="relu", padding="same")(pool2)

    deconv2 = Conv2DTranspose(128, 3, strides=2, padding="same")(convm)
    uconv2 = concatenate([deconv2, conv2])
    uconv2 = Conv2D(128, 3, activation="relu", padding="same")(uconv2)
    uconv2 = BatchNormalization()(uconv2)

    deconv1 = Conv2DTranspose(64, 3, strides=2, padding="same")(uconv2)
    uconv1 = concatenate([deconv1, conv1])
    uconv1 = Conv2D(128, 3, activation="relu", padding="same")(uconv1)
    uconv1 = BatchNormalization()(uconv1)

    out = Conv2D(1, 1, padding="same", activation="relu")(uconv1)

    model = Model(inp, out)

    #     decay_steps = EPOCHS * math.ceil(len(df) / BATCH_SIZE)
    #     scheduler = tf.keras.experimental.CosineDecay(
    #         initial_learning_rate=LR, decay_steps=decay_steps
    #     )
    #     optimizer = tf.keras.optimizers.Adam(learning_rate=scheduler)

    optimizer = tf.keras.optimizers.Adam(learning_rate=LR)
    loss = tf.keras.losses.LogCosh()

    model.compile(loss=loss, optimizer=optimizer, metrics=["mse"])

    return model<jupyter_output><empty_output><jupyter_text>Custom callback<jupyter_code>class CustomCallback(tf.keras.callbacks.Callback):
    """Early stopping based on CSI + MAE metrics"""

    def __init__(self, val_x, val_y, patience=5):
        super(CustomCallback, self).__init__()
        self.val_x = val_x
        self.val_y = val_y
        self.patience = patience
        self.best_weights = None

    def on_train_begin(self, logs=None):
        self.counter = 0
        self.best_score = np.inf
        self.stopped_epoch = 0

    def on_epoch_end(self, epoch, logs=None):
        preds = self.model.predict(self.val_x)
        y_pred = preds.reshape(len(preds), -1)
        y_true = self.val_y.reshape(len(self.val_y), -1)
        y_pred = y_pred[:, DAMS]
        y_true = y_true[:, DAMS]
        score = metrics.mean_absolute_error(y_true, y_pred)
        print("\nScore:", score)

        if np.less(score, self.best_score):
            self.best_score = score
            self.patience = 0
            self.best_weights = self.model.get_weights()
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.stopped_epoch = epoch
                self.model.stop_training = True
                print("\nRestoring model weights from the end of the best epoch.")
                self.model.set_weights(self.best_weights)

    def on_train_end(self, logs=None):
        if self.stopped_epoch > 0:
            print(f"Epoch {self.stopped_epoch:05d}: early stopping")<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>oof = []
targs = []

for fold in range(NUM_FOLDS):
    print("-" * 25)
    print(f"Training fold {fold}")

    train_df = df[df.fold != fold]
    val_df = df[df.fold == fold]

    train_gen = DataGenerator(train_df, batch_size=BATCH_SIZE, shuffle=True)

    # GET VALIDATION DATA
    val_gen = DataGenerator(val_df, batch_size=BATCH_SIZE * 4, shuffle=False)
    val_x, val_y = [], []
    for batch in val_gen:
        x, y = batch
        val_x.append(x)
        val_y.append(y)
    val_x = np.concatenate(val_x)
    val_y = np.concatenate(val_y)

    # MODEL
    with strategy.scope():
        model = get_baseline(INPUT_SHAPE)

    # CALLBACKS
#     early_stopping = tf.keras.callbacks.EarlyStopping(
#         monitor="mse", patience=3, mode="min", restore_best_weights=True
#     )
    scheduler = tf.keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.5, patience=5, verbose=1)
    custom_callback = CustomCallback(val_x, val_y, patience=5)
    history = model.fit(
        train_gen,
        epochs=EPOCHS,
        verbose=1,
        workers=3,
        callbacks=[custom_callback],
        validation_data=(val_x, val_y),
    )
    model.save_weights(f"baseline_fold{fold}.h5")

    oo = model.predict(val_x)
    oof.append(oo)
    targs.append(val_y)

    np.save(f"oof-fold{fold}", oo)
    np.save(f"val_y-fold{fold}", val_y)

    del model, train_df, val_df, train_gen, val_gen, val_x, val_y, oo, history, custom_callback
    K.clear_session()
    gc.collect()
    
oof = np.concatenate(oof)
targs = np.concatenate(targs)
oof = []
targs = []

for fold in range(NUM_FOLDS):
    oo = np.load(f"oof-fold{fold}.npy")
    oof.append(oo)
    targ = np.load(f"val_y-fold{fold}.npy")
    targs.append(targ)

oof = np.concatenate(oof)
targs = np.concatenate(targs)
oof.shape, targs.shape
y_true = targs.reshape(len(df), -1)
y_pred = oof.reshape(len(df), -1)
y_true = y_true[:, DAMS]
y_pred = y_pred[:, DAMS]
mae = metrics.mean_absolute_error(y_true, y_pred)
print(mae)<jupyter_output>7.452542545695364<jupyter_text>Inference<jupyter_code>def get_test():
    imgs = []
    test_paths = sorted(TEST_PATH.glob("*.npy"))
    for i, test_path in enumerate(test_paths):
        data = np.load(test_path)
        img = data[:, :, :4]
        img = np.where(img < 0, 0, img)
        #         img = img / 255
        #         img = img.astype(np.float32)
        imgs.append(img)

    imgs = np.stack(imgs)

    return imgs
test_imgs = get_test()
preds = np.zeros((len(test_imgs), SIZE * SIZE))

for fold in range(NUM_FOLDS):
    with strategy.scope():
        model = get_baseline(INPUT_SHAPE)
    model.load_weights(f"baseline_fold{fold}.h5")

    pred = model.predict(test_imgs, verbose=1)
    pred = pred.reshape(len(pred), -1)
    preds += pred

preds /= NUM_FOLDS
preds = np.round(preds)
preds = np.clip(preds, 0, 255)
preds = preds.astype(np.uint8)
np.save("preds", preds)
del test_imgs, preds, model
gc.collect()
preds = np.load("preds.npy")
file_names = [path.name for path in sorted(TEST_PATH.glob("*.npy"))]
subm = pd.DataFrame({"file_name": file_names})
for i in tqdm(range(SIZE * SIZE)):
    col = str(i)
    subm[col] = preds[:, i]
subm.to_csv("baseline2.csv", index=False)
subm.info()
subm.head()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14030939008769336, 'fraction_numerical': 0.020691693762932308, 'mean_word_length': 3.081222356252513, 'pattern_counts': {'":': 3, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '596821', 'license': 'MIT', 'n_tokens_mistral': 3974, 'n_tokens_neox': 3753, 'n_words': 680, 'path': 'notebooks/01-baseline.ipynb', 'repo_name': 'giwankim/precipitation-nowcasting', 'size': '10149'}