Spaces:
Sleeping
Sleeping
File size: 895 Bytes
b5cbaa6 8e6512c b5cbaa6 8e6512c b5cbaa6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
import segmentation_models_pytorch as smp
encoders = smp.encoders.encoders
WIDTH = 32
COLUMNS = [
"Encoder",
"Weights",
"Params, M",
]
def wrap_row(r):
return "|{}|".format(r)
header = "|".join([column.ljust(WIDTH, " ") for column in COLUMNS])
separator = "|".join(
["-" * WIDTH] + [":" + "-" * (WIDTH - 2) + ":"] * (len(COLUMNS) - 1)
)
print(wrap_row(header))
print(wrap_row(separator))
for encoder_name, encoder in encoders.items():
weights = "<br>".join(encoder["pretrained_settings"].keys())
encoder_name = encoder_name.ljust(WIDTH, " ")
weights = weights.ljust(WIDTH, " ")
model = encoder["encoder"](**encoder["params"], depth=5)
params = sum(p.numel() for p in model.parameters())
params = str(params // 1000000) + "M"
params = params.ljust(WIDTH, " ")
row = "|".join([encoder_name, weights, params])
print(wrap_row(row))
|