import torch
import torch.nn as nn

# https://jmlr.org/papers/v15/srivastava14a.html
# https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html

# Pravděpodobnost vynulování prvku (např. 0.5 = 50 %)
prob = 0.5

model = nn.Sequential(
    nn.Dropout(p=prob)
)
x = torch.ones(10)

# 1. Trénovací režim (výchozí stav model.train()):
# - Náhodně vynuluje prvky s pravděpodobností p.
# - Přeživší prvky vynásobí faktorem 1 / (1 - p) = 2, aby průměrná hodnota zůstala zachována (~1.0).
print("--- Trénovací režim (model.train()) ---")
print("Vstup: ", x)
output_train = model(x)
print("Výstup:", output_train)
print("Průměr:", torch.mean(output_train).item())

# 2. Testovací režim (model.eval()):
# - Dropout je zcela vypnutý (žádné nulování ani změna měřítka).
# - Výstup se přesně rovná vstupu.
print("\n--- Testovací režim (model.eval()) ---")
model.eval()
output_eval = model(x)
print("Výstup:", output_eval)
print("Průměr:", torch.mean(output_eval).item())