Holistically-Nested Edge Detection, retrained in TensorFlow
This notebook retrains HED (Xie & Tu, ICCV 2015) in TensorFlow/Keras, following the authors' original Caffe recipe, and exports weights that load directly into my hed-tensor library. It downloads its own data and runs end to end on Kaggle.
HED on Kaggle: train on T4 + export for hed-tensor¶
TensorFlow/Keras re-training of Holistically-Nested Edge Detection (Xie & Tu, ICCV 2015) following the authors' Caffe recipe (s9xie/hed). Exports weights that drop straight into the hed-tensor library. One notebook: it downloads its own data.
Runs on Kaggle with GPU T4 x2 and internet on. SMOKE_TEST = True does a quick end-to-end check before the full run.
Same as the Caffe recipe (solver.prototxt, train_val.prototxt): labels = round(gt/255) (only edges most annotators agree on) · class-balanced sigmoid cross-entropy summed over pixels for the 5 side outputs + fuse · SGD momentum 0.9, weight decay 2e-4 on weights only · per-layer learning rates (conv1-4 ×1, conv5 ×100, side outputs ×0.01, fuse ×0.001, biases ×2) · base lr 1e-6, ÷10 every 10k updates · 10 images per update · side outputs start at 0, fuse at 0.2.
Different: 'same' padding + bilinear resize instead of Caffe's pad=35 + deconvolution + crop (the layer layout of hed-tensor, so the weights stay compatible). The 10 images of an update are same-size images split over the GPUs instead of 10 accumulated single images: identical gradient, VGG has no batch norm.
Outputs (/kaggle/working): export/ (weights for the library), checkpoints/, logs/, figures/, predictions/*.zip.
Settings¶
import os
from pathlib import Path
SMOKE_TEST = False # True: tiny end-to-end run (a few minutes) to check everything before the real run
SMOKE_TEST = SMOKE_TEST or os.environ.get('HED_SMOKE_TEST') == '1' # (used for local testing)
# Training recipe
EPOCHS = 10 # 10 x 28,800 images ≈ the paper's 30k updates x 10 images
BATCH_SIZE = 10 # images per update (Caffe iter_size=10)
BASE_LR = 1e-6
LR_STEP_UPDATES = 10_000 # Caffe lr_policy "step": lr x LR_GAMMA every 10k updates
LR_GAMMA = 0.1
MOMENTUM = 0.9
WEIGHT_DECAY = 2e-4
LABEL_THRESHOLD = 128 # gt >= 128 is an edge, same as Caffe's round(gt/255)
SEED = 7
LOG_EVERY_UPDATES = 100
MAX_TRAIN_HOURS = 10.5 # checkpoint and move on to export before Kaggle's 12 h limit
if SMOKE_TEST:
EPOCHS, LOG_EVERY_UPDATES = 1, 2
SMOKE_TRAIN_IMAGES, SMOKE_TEST_IMAGES = 100, 6
DOWNLOADS = { # file name: (url, exact size in bytes)
'HED-BSDS.tar': ('https://vcl.ucsd.edu/hed/HED-BSDS.tar', 1_303_992_320),
'vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5': (
'https://storage.googleapis.com/tensorflow/keras-applications/vgg16/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5', 58_889_256),
}
HED_TENSOR = 'hed-tensor==1.0.1'
INPUT = Path(os.environ.get('HED_INPUT_DIR', '/kaggle/input')) # attached inputs: reused instead of downloading
WORK = Path(os.environ.get('HED_WORK_DIR', '/kaggle/working'))
DATA = Path(os.environ.get('HED_DATA_DIR', '/tmp/hed_data')) # downloads + unpacked data, not saved as output
PRED = Path(os.environ.get('HED_PRED_DIR', '/tmp/hed_predictions'))
EXPORT, CKPT_DIR, LOG_DIR, FIG_DIR = (WORK / d for d in ('export', 'checkpoints', 'logs', 'figures'))
for d in (WORK, DATA, PRED, EXPORT, CKPT_DIR, LOG_DIR, FIG_DIR):
d.mkdir(parents=True, exist_ok=True)
print('smoke test:', SMOKE_TEST)
smoke test: False
1. GPUs¶
Each update's 10 images are split in halves: on T4 x2 one half per GPU in parallel, on a single GPU the two halves run one after the other (keeps memory well inside 16 GB).
os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '1')
import time
import tensorflow as tf
from tensorflow import keras
build = tf.sysconfig.get_build_info()
print(f'TensorFlow {tf.__version__} | Keras {keras.__version__} | CUDA {build.get("cuda_version")} '
f'| built for {build.get("cuda_compute_capabilities")}')
gpus = tf.config.list_physical_devices('GPU')
for g in gpus:
tf.config.experimental.set_memory_growth(g, True)
GPU_NAMES = [tf.config.experimental.get_device_details(g).get('device_name', 'gpu') for g in gpus]
if len(gpus) >= 2:
DEVICES = ['/GPU:0', '/GPU:1']
elif gpus:
DEVICES = ['/GPU:0', '/GPU:0']
else:
if INPUT.exists() and not SMOKE_TEST:
raise SystemError('TensorFlow sees no GPU. Settings → Accelerator → GPU T4 x2.')
DEVICES = ['/CPU:0']
assert BATCH_SIZE % len(DEVICES) == 0
print('GPUs:', GPU_NAMES or 'none (CPU, only sensible for SMOKE_TEST)', '| each update split over', DEVICES)
for i, g in enumerate(gpus):
cc = tf.config.experimental.get_device_details(g).get('compute_capability')
if cc and cc[0] >= 10:
raise SystemError(f'{GPU_NAMES[i]} (compute capability {cc}) is a Blackwell GPU, which released TensorFlow '
'cannot run (CUDA_ERROR_INVALID_HANDLE). Switch the accelerator to GPU T4 x2.')
with tf.device(f'/GPU:{i}'):
t0 = time.time()
x = tf.random.normal([2, 320, 480, 64])
y = tf.image.resize(tf.nn.max_pool2d(tf.nn.relu(tf.nn.conv2d(x, tf.random.normal([3, 3, 64, 64]), 1, 'SAME')),
2, 2, 'SAME'), [320, 480])
float(tf.reduce_sum(y))
print(f'GPU {i} ({GPU_NAMES[i]}, compute capability {cc}) works: warm-up {time.time() - t0:.1f}s')
TensorFlow 2.20.0 | Keras 3.13.2 | CUDA 12.5.1 | built for ['sm_60', 'sm_70', 'sm_80', 'sm_89', 'compute_90'] GPUs: ['Tesla T4', 'Tesla T4'] | each update split over ['/GPU:0', '/GPU:1']
GPU 0 (Tesla T4, compute capability (7, 5)) works: warm-up 2.4s GPU 1 (Tesla T4, compute capability (7, 5)) works: warm-up 0.4s
2. Data, VGG16 weights, hed-tensor¶
Downloads the HED-BSDS training set (1.3 GB, UCSD) and the ImageNet VGG16 weights, resumable with retries and checked against their exact size. If the files are in an attached input (for example the old dependency-pack notebook), those are used instead. hed-tensor (your library) is installed to check the exported weights.
import subprocess
import sys
import urllib.request
def find_input(name):
hits = sorted(INPUT.rglob(name)) if INPUT.exists() else []
return hits[0] if hits else None
def _fetch(url, part):
"""Download url into part, resuming from what is already there."""
have = part.stat().st_size if part.exists() else 0
headers = {'Range': f'bytes={have}-'} if have else {}
with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=60) as r:
if have and r.status != 206: # server ignored the range request: start over
have = 0
total = have + int(r.headers.get('Content-Length') or 0)
done, next_report = have, have
with open(part, 'ab' if have else 'wb') as f:
while chunk := r.read(1 << 20):
f.write(chunk)
done += len(chunk)
if total > 50e6 and done >= next_report:
print(f' {part.name[:-5]}: {done / 1e6:,.0f} / {total / 1e6:,.0f} MB')
next_report = done + total // 10
def download(url, dest, size, tries=5):
if dest.exists() and dest.stat().st_size == size:
return dest
part = dest.with_name(dest.name + '.part')
for attempt in range(1, tries + 1):
try:
_fetch(url, part)
except Exception as e:
print(f' attempt {attempt}/{tries} failed: {e!r}')
got = part.stat().st_size if part.exists() else 0
if got == size:
break
if got > size: # corrupt partial file
part.unlink()
time.sleep(5 * attempt)
else:
raise RuntimeError(f'could not download {url}')
part.rename(dest)
return dest
def get_file(name):
url, size = DOWNLOADS[name]
found = find_input(name)
if found is not None and found.stat().st_size == size:
print(f'{name}: using attached input {found}')
return found
t0 = time.time()
path = download(url, DATA / name, size)
print(f'{name}: downloaded ({size / 1e6:,.0f} MB, {time.time() - t0:.0f}s)')
return path
if not list(DATA.rglob('train_pair.lst')):
tar = get_file('HED-BSDS.tar')
t0 = time.time()
subprocess.run(['tar', '-xf', str(tar), '-C', str(DATA)], check=True)
print(f'unpacked HED-BSDS.tar in {time.time() - t0:.0f}s')
HED_ROOT = next(DATA.rglob('train_pair.lst')).parent
VGG_WEIGHTS = get_file('vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5')
print('data:', HED_ROOT)
try:
import hed_tensor
except ImportError:
cmd = [sys.executable, '-m', 'pip', 'install', '-q', '--no-deps']
r = subprocess.run(cmd + [HED_TENSOR], capture_output=True, text=True)
if r.returncode and find_input('hed_tensor-1.0.1-py3-none-any.whl'): # no internet: wheel from an attached input
r = subprocess.run(cmd + [str(find_input('hed_tensor-1.0.1-py3-none-any.whl'))], capture_output=True, text=True)
if r.returncode:
raise RuntimeError(r.stderr)
import hed_tensor
print('hed_tensor', hed_tensor.__version__)
HED-BSDS.tar: 1 / 1,304 MB
HED-BSDS.tar: 132 / 1,304 MB
HED-BSDS.tar: 263 / 1,304 MB
HED-BSDS.tar: 394 / 1,304 MB
HED-BSDS.tar: 525 / 1,304 MB
HED-BSDS.tar: 656 / 1,304 MB
HED-BSDS.tar: 787 / 1,304 MB
HED-BSDS.tar: 919 / 1,304 MB
HED-BSDS.tar: 1,050 / 1,304 MB
HED-BSDS.tar: 1,181 / 1,304 MB
HED-BSDS.tar: downloaded (1,304 MB, 134s)
unpacked HED-BSDS.tar in 4s
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 1 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 7 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 14 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 20 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 26 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 33 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 39 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 45 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 51 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: 58 / 59 MB
vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5: downloaded (59 MB, 1s)
data: /tmp/hed_data/HED-BSDS
hed_tensor 1.0.0
3. Data pipeline¶
The 28,800 augmented training pairs come in only 27 image sizes, so each batch is filled with images of one size: no padding or resizing. Batches are reshuffled every epoch. The 200 BSDS500 test images (test/ in HED-BSDS) are used for progress figures and predictions.
import random
from collections import defaultdict
import numpy as np
from PIL import Image
MEAN_BGR = np.array([104.00699, 116.66877, 122.67892], np.float32) # Caffe HED / VGG mean, BGR order
pairs = [l.split() for l in (HED_ROOT / 'train_pair.lst').read_text().splitlines() if l.strip()]
if SMOKE_TEST:
pairs = pairs[:SMOKE_TRAIN_IMAGES] # few image sizes, so full batches of BATCH_SIZE
IMG_PATHS = np.array([str(HED_ROOT / img) for img, _ in pairs])
LBL_PATHS = np.array([str(HED_ROOT / lbl) for _, lbl in pairs])
t0 = time.time()
by_size = defaultdict(list)
for i, path in enumerate(IMG_PATHS):
by_size[Image.open(path).size].append(i) # reads the header only
print(f'{len(pairs):,} training pairs in {len(by_size)} image sizes ({time.time() - t0:.0f}s)')
def epoch_order(epoch):
"""Image indices for one epoch; every run of BATCH_SIZE consecutive indices has one image size."""
rng = random.Random(SEED * 1000 + epoch)
batches = []
for idx in by_size.values():
idx = idx[:]
rng.shuffle(idx)
batches += [idx[i:i + BATCH_SIZE] for i in range(0, len(idx) - BATCH_SIZE + 1, BATCH_SIZE)]
rng.shuffle(batches)
return [i for b in batches for i in b]
UPDATES_PER_EPOCH = len(epoch_order(0)) // BATCH_SIZE
TOTAL_UPDATES = EPOCHS * UPDATES_PER_EPOCH
def load_pair(img_path, lbl_path):
img = tf.io.decode_image(tf.io.read_file(img_path), channels=3, expand_animations=False)
img = tf.cast(img[..., ::-1], tf.float32) - MEAN_BGR # RGB → BGR, minus mean
gt = tf.io.decode_image(tf.io.read_file(lbl_path), channels=1, expand_animations=False)
return img, tf.cast(gt >= LABEL_THRESHOLD, tf.float32)
def make_dataset(epoch):
order = epoch_order(epoch)
ds = tf.data.Dataset.from_tensor_slices((IMG_PATHS[order], LBL_PATHS[order]))
ds = ds.map(load_pair, num_parallel_calls=tf.data.AUTOTUNE, deterministic=True)
return ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
TEST_PATHS = [HED_ROOT / l.split()[0] for l in (HED_ROOT / 'test.lst').read_text().splitlines() if l.strip()]
if SMOKE_TEST:
TEST_PATHS = TEST_PATHS[:SMOKE_TEST_IMAGES]
def load_test(path):
"""RGB uint8 image and the model input (1, H, W, 3)."""
rgb = np.array(Image.open(path).convert('RGB'))
return rgb, (rgb[..., ::-1].astype(np.float32) - MEAN_BGR)[None]
images, labels = next(iter(make_dataset(0)))
print(f'{UPDATES_PER_EPOCH:,} updates per epoch, {TOTAL_UPDATES:,} total | first batch {images.shape}, '
f'edge pixels {float(tf.reduce_mean(labels)):.1%} | {len(TEST_PATHS)} test images')
28,800 training pairs in 27 image sizes (2s) 2,868 updates per epoch, 28,680 total | first batch (10, 192, 105, 3), edge pixels 0.1% | 200 test images
4. Model¶
Same attribute names and layer order as hed_tensor.HED (block1..5, dsn1..5, fuse), so model.save_weights() gives a file the library loads unchanged. Only the starting values follow Caffe.
from tensorflow.keras import layers
class HED(keras.Model):
def __init__(self, vgg_weights=None):
super().__init__()
vgg = keras.applications.VGG16(weights=vgg_weights, include_top=False)
get = vgg.get_layer
def pool():
return layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same')
self.block1 = keras.Sequential([get('block1_conv1'), get('block1_conv2')])
self.block2 = keras.Sequential([pool(), get('block2_conv1'), get('block2_conv2')])
self.block3 = keras.Sequential([pool(), get('block3_conv1'), get('block3_conv2'), get('block3_conv3')])
self.block4 = keras.Sequential([pool(), get('block4_conv1'), get('block4_conv2'), get('block4_conv3')])
self.block5 = keras.Sequential([pool(), get('block5_conv1'), get('block5_conv2'), get('block5_conv3')])
# Caffe: side outputs have no filler (start at 0), the fuse layer starts at 0.2
self.dsn1, self.dsn2, self.dsn3, self.dsn4, self.dsn5 = (
layers.Conv2D(1, 1, padding='same', kernel_initializer='zeros') for _ in range(5))
self.fuse = layers.Conv2D(1, 1, padding='same', kernel_initializer=keras.initializers.Constant(0.2))
def call(self, x, training=False):
"""training=True: logits of the 5 side outputs + fuse. training=False: sigmoid(fuse)."""
h, w = tf.shape(x)[1], tf.shape(x)[2]
c1 = self.block1(x)
c2 = self.block2(c1)
c3 = self.block3(c2)
c4 = self.block4(c3)
c5 = self.block5(c4)
sides = [self.dsn1(c1)] + [tf.image.resize(dsn(c), [h, w], method='bilinear')
for dsn, c in ((self.dsn2, c2), (self.dsn3, c3), (self.dsn4, c4), (self.dsn5, c5))]
fused = self.fuse(tf.concat(sides, axis=-1))
return (*sides, fused) if training else tf.nn.sigmoid(fused)
def blocks(self):
return [self.block1, self.block2, self.block3, self.block4, self.block5]
def dsns(self):
return [self.dsn1, self.dsn2, self.dsn3, self.dsn4, self.dsn5]
model = HED(vgg_weights=str(VGG_WEIGHTS))
model(tf.zeros([1, 64, 64, 3]))
print(f'{sum(int(np.prod(v.shape)) for v in model.trainable_variables) / 1e6:.2f}M parameters')
14.72M parameters
5. Loss and Caffe-style SGD¶
Loss per output = Caffe's HED SigmoidCrossEntropyLoss: edge pixels weighted by the fraction of non-edge pixels and vice versa, summed over the image, averaged over the images in the update. Total = 5 side outputs + fuse.
Update (Caffe SGD, per parameter): v = momentum·v + lr·lr_mult·(grad + weight_decay·decay_mult·w), w -= v.
The training log also shows each loss relative to chance (an output that always says 0.5): 1.0 = guessing, lower is better. Raw summed losses jump around with the number of edge pixels per batch; the ratio doesn't. Expect a short spike above 1 in the first updates (momentum + conv5's 100x learning rate), as in Caffe.
import math
def balanced_bce(labels, logits):
"""Per-image Caffe HED loss, shape (batch,)."""
bce = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)
n_all = tf.cast(tf.reduce_prod(tf.shape(labels)[1:]), tf.float32)
n_pos = tf.reduce_sum(labels, axis=[1, 2, 3])
loss_pos = tf.reduce_sum(labels * bce, axis=[1, 2, 3])
loss_neg = tf.reduce_sum((1.0 - labels) * bce, axis=[1, 2, 3])
return loss_pos * (n_all - n_pos) / n_all + loss_neg * n_pos / n_all
def chance_loss(labels):
"""Per-image balanced_bce of an output that always says 0.5."""
n_all = tf.cast(tf.reduce_prod(tf.shape(labels)[1:]), tf.float32)
n_pos = tf.reduce_sum(labels, axis=[1, 2, 3])
return 2.0 * math.log(2.0) * n_pos * (n_all - n_pos) / n_all
def caffe_param_groups(model):
"""(variable, lr_mult, decay_mult) as in s9xie/hed train_val.prototxt."""
convs = [(l, 100.0 if blk is model.block5 else 1.0)
for blk in model.blocks() for l in blk.layers if isinstance(l, layers.Conv2D)]
convs += [(dsn, 0.01) for dsn in model.dsns()] + [(model.fuse, 0.001)]
groups = [g for layer, m in convs for g in ((layer.kernel, m, 1.0), (layer.bias, 2 * m, 0.0))]
assert len(groups) == len(model.trainable_variables)
return groups
VARS, LR_MULTS, DECAY_MULTS = zip(*caffe_param_groups(model))
MOMENTA = [tf.Variable(tf.zeros_like(v), trainable=False) for v in VARS]
UPDATE = tf.Variable(0, dtype=tf.int64, trainable=False) # updates done (saved in checkpoints)
EPOCH = tf.Variable(0, dtype=tf.int64, trainable=False) # epochs done
@tf.function(input_signature=[tf.TensorSpec([None, None, None, 3], tf.float32),
tf.TensorSpec([None, None, None, 1], tf.float32),
tf.TensorSpec([], tf.float32)])
def train_step(images, labels, lr):
"""One Caffe SGD update on the whole batch, computed in len(DEVICES) parts."""
n = tf.cast(tf.shape(images)[0], tf.float32)
losses, chance, grads, prev = 0.0, 0.0, None, None
for k, (dev, x, y) in enumerate(zip(DEVICES, tf.split(images, len(DEVICES)), tf.split(labels, len(DEVICES)))):
# parts on the same GPU run one after the other, so only one part's activations are in memory
with tf.device(dev), tf.control_dependencies(grads if k and dev == prev else []):
with tf.GradientTape() as tape:
part = tf.stack([tf.reduce_sum(balanced_bce(y, out)) for out in model(x, training=True)]) / n
total = tf.reduce_sum(part)
part_grads = tape.gradient(total, VARS)
grads = part_grads if grads is None else [a + b for a, b in zip(grads, part_grads)]
losses += part
chance += tf.reduce_sum(chance_loss(y)) / n
prev = dev
for v, g, m, lr_mult, decay_mult in zip(VARS, grads, MOMENTA, LR_MULTS, DECAY_MULTS):
if decay_mult:
g = g + WEIGHT_DECAY * decay_mult * v
m.assign(MOMENTUM * m + lr * lr_mult * g)
v.assign_sub(m)
return losses, chance
@tf.function(input_signature=[tf.TensorSpec([1, None, None, 3], tf.float32)])
def predict_all(x):
"""Sigmoid maps of the 5 side outputs + fuse for one image."""
return [tf.nn.sigmoid(o) for o in model(x, training=True)]
print(f'lr multipliers: {sorted(set(LR_MULTS))}')
lr multipliers: [0.001, 0.002, 0.01, 0.02, 1.0, 2.0, 100.0, 200.0]
6. Train¶
Checkpoint (weights + momentum + counters) after every epoch, kept in checkpoints/. A progress figure on 3 test images is saved per epoch.
import csv
import shutil
import matplotlib.pyplot as plt
if not (CKPT_DIR / 'checkpoint').exists() and find_input('checkpoint') is not None:
prev_run = find_input('checkpoint').parent.parent # .../checkpoints/checkpoint of an attached earlier version
shutil.copytree(prev_run / 'checkpoints', CKPT_DIR, dirs_exist_ok=True)
if (prev_run / 'logs/train_log.csv').exists():
shutil.copy(prev_run / 'logs/train_log.csv', LOG_DIR)
print('copied checkpoint from attached input', prev_run)
ckpt = tf.train.Checkpoint(model=model, momenta=MOMENTA, update=UPDATE, epoch=EPOCH)
manager = tf.train.CheckpointManager(ckpt, CKPT_DIR, max_to_keep=2)
if manager.latest_checkpoint:
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()
print(f'resumed {manager.latest_checkpoint}: {int(EPOCH.numpy())} epochs, {int(UPDATE.numpy()):,} updates done')
else:
print('starting from ImageNet VGG16')
PREVIEW = TEST_PATHS[:3]
def save_progress_figure(epoch):
fig, axes = plt.subplots(2, len(PREVIEW), figsize=(4 * len(PREVIEW), 6))
for col, path in enumerate(PREVIEW):
rgb, x = load_test(path)
axes[0, col].imshow(rgb)
axes[1, col].imshow(1 - predict_all(x)[-1].numpy()[0, ..., 0], cmap='gray', vmin=0, vmax=1)
for ax in axes.flat:
ax.axis('off')
fig.suptitle(f'fused output after epoch {epoch}')
fig.tight_layout()
fig.savefig(FIG_DIR / f'progress_epoch_{epoch:02d}.png', dpi=80)
plt.close(fig)
LOG_PATH = LOG_DIR / 'train_log.csv'
if not LOG_PATH.exists():
LOG_PATH.write_text('update,epoch,lr,loss,dsn1,dsn2,dsn3,dsn4,dsn5,fuse,chance,img_per_s\n')
update, start_epoch = int(UPDATE.numpy()), int(EPOCH.numpy())
session_start, stopped_early = time.time(), False
print(f'training epochs {start_epoch + 1}..{EPOCHS}, {TOTAL_UPDATES - update:,} updates to go')
with open(LOG_PATH, 'a', newline='') as log_file:
log = csv.writer(log_file)
for epoch in range(start_epoch, EPOCHS):
window, n_images, t_window = [], 0, time.time()
for b, (images, labels) in enumerate(make_dataset(epoch)):
lr = BASE_LR * LR_GAMMA ** (update // LR_STEP_UPDATES)
window.append(train_step(images, labels, tf.constant(lr, tf.float32)))
update += 1
n_images += int(images.shape[0])
if update % LOG_EVERY_UPDATES:
continue
losses = tf.reduce_mean(tf.stack([l for l, _ in window]), axis=0).numpy()
chance = float(tf.reduce_mean(tf.stack([c for _, c in window])))
if not np.isfinite(losses).all():
raise FloatingPointError(f'loss became {losses} at update {update}')
speed = n_images / (time.time() - t_window)
eta_min = (TOTAL_UPDATES - update) * BATCH_SIZE / speed / 60
log.writerow([update, round(epoch + (b + 1) / UPDATES_PER_EPOCH, 4), lr, losses.sum(), *losses, chance,
round(speed, 1)])
log_file.flush()
print(f'epoch {epoch + 1}/{EPOCHS} update {update:,}/{TOTAL_UPDATES:,} loss {losses.sum():,.0f} '
f'vs chance: fuse {losses[-1] / chance:.2f}x, sides {losses[:5].mean() / chance:.2f}x '
f'lr {lr:.0e} {speed:.0f} img/s ETA {eta_min:.0f} min')
window, n_images, t_window = [], 0, time.time()
if time.time() - session_start > MAX_TRAIN_HOURS * 3600:
stopped_early = True
break
UPDATE.assign(update)
if stopped_early:
manager.save()
print(f'Stopped after {MAX_TRAIN_HOURS} h, checkpoint saved. Attach this version as input and run again to finish '
f'(the unfinished epoch restarts). Continuing with export of the current weights.')
break
EPOCH.assign(epoch + 1)
manager.save()
save_progress_figure(epoch + 1)
print(f'--- epoch {epoch + 1} done, checkpoint saved ({(time.time() - session_start) / 60:.0f} min this session)')
TRAINING_COMPLETE = int(EPOCH.numpy()) >= EPOCHS
print('training complete' if TRAINING_COMPLETE else 'training NOT complete')
starting from ImageNet VGG16 training epochs 1..10, 28,680 updates to go epoch 1/10 update 100/28,680 loss 10,440 vs chance: fuse 2.06x, sides 2.54x lr 1e-06 6 img/s ETA 795 min epoch 1/10 update 200/28,680 loss 4,419 vs chance: fuse 0.85x, sides 1.01x lr 1e-06 33 img/s ETA 143 min epoch 1/10 update 300/28,680 loss 3,214 vs chance: fuse 0.82x, sides 0.92x lr 1e-06 29 img/s ETA 164 min epoch 1/10 update 400/28,680 loss 3,378 vs chance: fuse 0.77x, sides 0.86x lr 1e-06 38 img/s ETA 125 min epoch 1/10 update 500/28,680 loss 3,016 vs chance: fuse 0.76x, sides 0.85x lr 1e-06 45 img/s ETA 105 min epoch 1/10 update 600/28,680 loss 2,882 vs chance: fuse 0.74x, sides 0.88x lr 1e-06 48 img/s ETA 98 min epoch 1/10 update 700/28,680 loss 3,325 vs chance: fuse 0.73x, sides 0.82x lr 1e-06 38 img/s ETA 124 min epoch 1/10 update 800/28,680 loss 2,923 vs chance: fuse 0.79x, sides 0.88x lr 1e-06 44 img/s ETA 105 min epoch 1/10 update 900/28,680 loss 3,319 vs chance: fuse 0.74x, sides 0.84x lr 1e-06 43 img/s ETA 108 min epoch 1/10 update 1,000/28,680 loss 2,494 vs chance: fuse 0.72x, sides 0.81x lr 1e-06 43 img/s ETA 106 min epoch 1/10 update 1,100/28,680 loss 2,873 vs chance: fuse 0.70x, sides 0.77x lr 1e-06 39 img/s ETA 119 min epoch 1/10 update 1,200/28,680 loss 2,920 vs chance: fuse 0.67x, sides 0.75x lr 1e-06 38 img/s ETA 121 min epoch 1/10 update 1,300/28,680 loss 3,266 vs chance: fuse 0.66x, sides 0.75x lr 1e-06 35 img/s ETA 130 min epoch 1/10 update 1,400/28,680 loss 2,569 vs chance: fuse 0.69x, sides 0.79x lr 1e-06 42 img/s ETA 107 min epoch 1/10 update 1,500/28,680 loss 2,718 vs chance: fuse 0.66x, sides 0.74x lr 1e-06 40 img/s ETA 115 min epoch 1/10 update 1,600/28,680 loss 2,888 vs chance: fuse 0.67x, sides 0.75x lr 1e-06 38 img/s ETA 118 min epoch 1/10 update 1,700/28,680 loss 2,332 vs chance: fuse 0.65x, sides 0.73x lr 1e-06 41 img/s ETA 108 min epoch 1/10 update 1,800/28,680 loss 2,554 vs chance: fuse 0.65x, sides 0.72x lr 1e-06 41 img/s ETA 110 min epoch 1/10 update 1,900/28,680 loss 2,205 vs chance: fuse 0.65x, sides 0.73x lr 1e-06 46 img/s ETA 97 min epoch 1/10 update 2,000/28,680 loss 2,875 vs chance: fuse 0.63x, sides 0.71x lr 1e-06 36 img/s ETA 122 min epoch 1/10 update 2,100/28,680 loss 2,895 vs chance: fuse 0.65x, sides 0.73x lr 1e-06 40 img/s ETA 112 min epoch 1/10 update 2,200/28,680 loss 2,976 vs chance: fuse 0.62x, sides 0.70x lr 1e-06 35 img/s ETA 128 min epoch 1/10 update 2,300/28,680 loss 2,267 vs chance: fuse 0.63x, sides 0.71x lr 1e-06 42 img/s ETA 106 min epoch 1/10 update 2,400/28,680 loss 3,001 vs chance: fuse 0.62x, sides 0.70x lr 1e-06 37 img/s ETA 117 min epoch 1/10 update 2,500/28,680 loss 2,609 vs chance: fuse 0.61x, sides 0.69x lr 1e-06 38 img/s ETA 116 min epoch 1/10 update 2,600/28,680 loss 2,750 vs chance: fuse 0.63x, sides 0.71x lr 1e-06 41 img/s ETA 105 min epoch 1/10 update 2,700/28,680 loss 2,831 vs chance: fuse 0.62x, sides 0.70x lr 1e-06 35 img/s ETA 125 min epoch 1/10 update 2,800/28,680 loss 2,600 vs chance: fuse 0.61x, sides 0.70x lr 1e-06 43 img/s ETA 101 min --- epoch 1 done, checkpoint saved (15 min this session) epoch 2/10 update 2,900/28,680 loss 2,958 vs chance: fuse 0.59x, sides 0.68x lr 1e-06 34 img/s ETA 125 min epoch 2/10 update 3,000/28,680 loss 2,651 vs chance: fuse 0.62x, sides 0.71x lr 1e-06 42 img/s ETA 101 min epoch 2/10 update 3,100/28,680 loss 2,476 vs chance: fuse 0.60x, sides 0.68x lr 1e-06 41 img/s ETA 105 min epoch 2/10 update 3,200/28,680 loss 2,726 vs chance: fuse 0.60x, sides 0.68x lr 1e-06 37 img/s ETA 116 min epoch 2/10 update 3,300/28,680 loss 2,908 vs chance: fuse 0.59x, sides 0.68x lr 1e-06 33 img/s ETA 127 min epoch 2/10 update 3,400/28,680 loss 2,306 vs chance: fuse 0.59x, sides 0.67x lr 1e-06 42 img/s ETA 101 min epoch 2/10 update 3,500/28,680 loss 2,463 vs chance: fuse 0.58x, sides 0.67x lr 1e-06 41 img/s ETA 101 min epoch 2/10 update 3,600/28,680 loss 2,791 vs chance: fuse 0.59x, sides 0.68x lr 1e-06 38 img/s ETA 111 min epoch 2/10 update 3,700/28,680 loss 2,187 vs chance: fuse 0.58x, sides 0.67x lr 1e-06 41 img/s ETA 101 min epoch 2/10 update 3,800/28,680 loss 1,897 vs chance: fuse 0.58x, sides 0.67x lr 1e-06 47 img/s ETA 88 min epoch 2/10 update 3,900/28,680 loss 2,450 vs chance: fuse 0.59x, sides 0.67x lr 1e-06 40 img/s ETA 103 min epoch 2/10 update 4,000/28,680 loss 2,367 vs chance: fuse 0.58x, sides 0.67x lr 1e-06 39 img/s ETA 105 min epoch 2/10 update 4,100/28,680 loss 2,439 vs chance: fuse 0.56x, sides 0.65x lr 1e-06 39 img/s ETA 106 min epoch 2/10 update 4,200/28,680 loss 2,604 vs chance: fuse 0.59x, sides 0.68x lr 1e-06 39 img/s ETA 105 min epoch 2/10 update 4,300/28,680 loss 2,254 vs chance: fuse 0.57x, sides 0.66x lr 1e-06 41 img/s ETA 100 min epoch 2/10 update 4,400/28,680 loss 2,604 vs chance: fuse 0.58x, sides 0.67x lr 1e-06 42 img/s ETA 97 min epoch 2/10 update 4,500/28,680 loss 3,014 vs chance: fuse 0.58x, sides 0.67x lr 1e-06 35 img/s ETA 114 min epoch 2/10 update 4,600/28,680 loss 2,460 vs chance: fuse 0.57x, sides 0.66x lr 1e-06 38 img/s ETA 106 min epoch 2/10 update 4,700/28,680 loss 2,696 vs chance: fuse 0.57x, sides 0.66x lr 1e-06 38 img/s ETA 105 min epoch 2/10 update 4,800/28,680 loss 2,578 vs chance: fuse 0.57x, sides 0.67x lr 1e-06 39 img/s ETA 103 min epoch 2/10 update 4,900/28,680 loss 2,564 vs chance: fuse 0.55x, sides 0.65x lr 1e-06 37 img/s ETA 107 min epoch 2/10 update 5,000/28,680 loss 2,262 vs chance: fuse 0.56x, sides 0.66x lr 1e-06 43 img/s ETA 93 min epoch 2/10 update 5,100/28,680 loss 2,033 vs chance: fuse 0.56x, sides 0.65x lr 1e-06 44 img/s ETA 89 min epoch 2/10 update 5,200/28,680 loss 2,298 vs chance: fuse 0.55x, sides 0.64x lr 1e-06 36 img/s ETA 108 min epoch 2/10 update 5,300/28,680 loss 3,080 vs chance: fuse 0.56x, sides 0.65x lr 1e-06 33 img/s ETA 117 min epoch 2/10 update 5,400/28,680 loss 1,995 vs chance: fuse 0.56x, sides 0.65x lr 1e-06 43 img/s ETA 90 min epoch 2/10 update 5,500/28,680 loss 1,926 vs chance: fuse 0.56x, sides 0.65x lr 1e-06 48 img/s ETA 81 min epoch 2/10 update 5,600/28,680 loss 2,605 vs chance: fuse 0.55x, sides 0.64x lr 1e-06 38 img/s ETA 100 min epoch 2/10 update 5,700/28,680 loss 2,427 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 38 img/s ETA 100 min --- epoch 2 done, checkpoint saved (27 min this session) epoch 3/10 update 5,800/28,680 loss 2,067 vs chance: fuse 0.56x, sides 0.66x lr 1e-06 42 img/s ETA 92 min epoch 3/10 update 5,900/28,680 loss 2,493 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 38 img/s ETA 100 min epoch 3/10 update 6,000/28,680 loss 2,104 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 41 img/s ETA 91 min epoch 3/10 update 6,100/28,680 loss 2,039 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 44 img/s ETA 85 min epoch 3/10 update 6,200/28,680 loss 2,490 vs chance: fuse 0.53x, sides 0.63x lr 1e-06 37 img/s ETA 101 min epoch 3/10 update 6,300/28,680 loss 2,200 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 39 img/s ETA 95 min epoch 3/10 update 6,400/28,680 loss 2,411 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 37 img/s ETA 101 min epoch 3/10 update 6,500/28,680 loss 2,225 vs chance: fuse 0.53x, sides 0.64x lr 1e-06 40 img/s ETA 92 min epoch 3/10 update 6,600/28,680 loss 2,449 vs chance: fuse 0.54x, sides 0.63x lr 1e-06 40 img/s ETA 93 min epoch 3/10 update 6,700/28,680 loss 2,485 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 37 img/s ETA 99 min epoch 3/10 update 6,800/28,680 loss 2,435 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 42 img/s ETA 87 min epoch 3/10 update 6,900/28,680 loss 2,548 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 37 img/s ETA 98 min epoch 3/10 update 7,000/28,680 loss 2,462 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 39 img/s ETA 94 min epoch 3/10 update 7,100/28,680 loss 2,405 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 40 img/s ETA 90 min epoch 3/10 update 7,200/28,680 loss 2,278 vs chance: fuse 0.53x, sides 0.63x lr 1e-06 37 img/s ETA 96 min epoch 3/10 update 7,300/28,680 loss 2,190 vs chance: fuse 0.54x, sides 0.64x lr 1e-06 44 img/s ETA 81 min epoch 3/10 update 7,400/28,680 loss 2,438 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 36 img/s ETA 97 min epoch 3/10 update 7,500/28,680 loss 2,261 vs chance: fuse 0.52x, sides 0.63x lr 1e-06 39 img/s ETA 91 min epoch 3/10 update 7,600/28,680 loss 2,532 vs chance: fuse 0.52x, sides 0.63x lr 1e-06 35 img/s ETA 99 min epoch 3/10 update 7,700/28,680 loss 2,523 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 35 img/s ETA 100 min epoch 3/10 update 7,800/28,680 loss 2,651 vs chance: fuse 0.53x, sides 0.63x lr 1e-06 36 img/s ETA 95 min epoch 3/10 update 7,900/28,680 loss 2,162 vs chance: fuse 0.53x, sides 0.63x lr 1e-06 43 img/s ETA 80 min epoch 3/10 update 8,000/28,680 loss 1,803 vs chance: fuse 0.52x, sides 0.63x lr 1e-06 48 img/s ETA 72 min epoch 3/10 update 8,100/28,680 loss 2,361 vs chance: fuse 0.52x, sides 0.63x lr 1e-06 41 img/s ETA 84 min epoch 3/10 update 8,200/28,680 loss 2,742 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 35 img/s ETA 96 min epoch 3/10 update 8,300/28,680 loss 2,161 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 40 img/s ETA 85 min epoch 3/10 update 8,400/28,680 loss 2,103 vs chance: fuse 0.52x, sides 0.63x lr 1e-06 41 img/s ETA 82 min epoch 3/10 update 8,500/28,680 loss 2,579 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 37 img/s ETA 91 min epoch 3/10 update 8,600/28,680 loss 2,077 vs chance: fuse 0.53x, sides 0.63x lr 1e-06 44 img/s ETA 77 min --- epoch 3 done, checkpoint saved (39 min this session) epoch 4/10 update 8,700/28,680 loss 2,223 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 39 img/s ETA 85 min epoch 4/10 update 8,800/28,680 loss 2,926 vs chance: fuse 0.50x, sides 0.61x lr 1e-06 35 img/s ETA 95 min epoch 4/10 update 8,900/28,680 loss 2,590 vs chance: fuse 0.52x, sides 0.62x lr 1e-06 36 img/s ETA 91 min epoch 4/10 update 9,000/28,680 loss 1,935 vs chance: fuse 0.51x, sides 0.62x lr 1e-06 42 img/s ETA 79 min epoch 4/10 update 9,100/28,680 loss 2,415 vs chance: fuse 0.53x, sides 0.64x lr 1e-06 39 img/s ETA 84 min epoch 4/10 update 9,200/28,680 loss 2,407 vs chance: fuse 0.51x, sides 0.62x lr 1e-06 38 img/s ETA 86 min epoch 4/10 update 9,300/28,680 loss 2,143 vs chance: fuse 0.52x, sides 0.63x lr 1e-06 42 img/s ETA 77 min epoch 4/10 update 9,400/28,680 loss 2,467 vs chance: fuse 0.51x, sides 0.62x lr 1e-06 37 img/s ETA 87 min epoch 4/10 update 9,500/28,680 loss 2,003 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 43 img/s ETA 75 min epoch 4/10 update 9,600/28,680 loss 2,558 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 36 img/s ETA 87 min epoch 4/10 update 9,700/28,680 loss 2,138 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 41 img/s ETA 77 min epoch 4/10 update 9,800/28,680 loss 2,075 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 40 img/s ETA 79 min epoch 4/10 update 9,900/28,680 loss 2,254 vs chance: fuse 0.51x, sides 0.61x lr 1e-06 41 img/s ETA 76 min epoch 4/10 update 10,000/28,680 loss 2,048 vs chance: fuse 0.50x, sides 0.61x lr 1e-06 41 img/s ETA 76 min epoch 4/10 update 10,100/28,680 loss 2,129 vs chance: fuse 0.51x, sides 0.61x lr 1e-07 41 img/s ETA 75 min epoch 4/10 update 10,200/28,680 loss 2,496 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 37 img/s ETA 83 min epoch 4/10 update 10,300/28,680 loss 2,304 vs chance: fuse 0.50x, sides 0.60x lr 1e-07 38 img/s ETA 80 min epoch 4/10 update 10,400/28,680 loss 2,306 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 38 img/s ETA 80 min epoch 4/10 update 10,500/28,680 loss 1,756 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 46 img/s ETA 67 min epoch 4/10 update 10,600/28,680 loss 2,019 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 43 img/s ETA 70 min epoch 4/10 update 10,700/28,680 loss 2,467 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 36 img/s ETA 83 min epoch 4/10 update 10,800/28,680 loss 2,397 vs chance: fuse 0.46x, sides 0.57x lr 1e-07 34 img/s ETA 87 min epoch 4/10 update 10,900/28,680 loss 1,867 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 74 min epoch 4/10 update 11,000/28,680 loss 2,119 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 41 img/s ETA 73 min epoch 4/10 update 11,100/28,680 loss 1,874 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 44 img/s ETA 66 min epoch 4/10 update 11,200/28,680 loss 2,369 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 35 img/s ETA 83 min epoch 4/10 update 11,300/28,680 loss 2,131 vs chance: fuse 0.49x, sides 0.59x lr 1e-07 41 img/s ETA 71 min epoch 4/10 update 11,400/28,680 loss 2,055 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 72 min --- epoch 4 done, checkpoint saved (51 min this session) epoch 5/10 update 11,500/28,680 loss 2,089 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 43 img/s ETA 66 min epoch 5/10 update 11,600/28,680 loss 2,359 vs chance: fuse 0.49x, sides 0.59x lr 1e-07 37 img/s ETA 78 min epoch 5/10 update 11,700/28,680 loss 2,016 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 39 img/s ETA 72 min epoch 5/10 update 11,800/28,680 loss 2,477 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 40 img/s ETA 71 min epoch 5/10 update 11,900/28,680 loss 1,940 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 44 img/s ETA 63 min epoch 5/10 update 12,000/28,680 loss 2,412 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 37 img/s ETA 74 min epoch 5/10 update 12,100/28,680 loss 2,146 vs chance: fuse 0.48x, sides 0.58x lr 1e-07 39 img/s ETA 71 min epoch 5/10 update 12,200/28,680 loss 2,501 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 37 img/s ETA 75 min epoch 5/10 update 12,300/28,680 loss 2,019 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 68 min epoch 5/10 update 12,400/28,680 loss 2,214 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 68 min epoch 5/10 update 12,500/28,680 loss 2,418 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 36 img/s ETA 74 min epoch 5/10 update 12,600/28,680 loss 2,142 vs chance: fuse 0.50x, sides 0.61x lr 1e-07 42 img/s ETA 64 min epoch 5/10 update 12,700/28,680 loss 1,901 vs chance: fuse 0.49x, sides 0.59x lr 1e-07 42 img/s ETA 64 min epoch 5/10 update 12,800/28,680 loss 1,821 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 39 img/s ETA 68 min epoch 5/10 update 12,900/28,680 loss 2,086 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 41 img/s ETA 64 min epoch 5/10 update 13,000/28,680 loss 2,151 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 39 img/s ETA 68 min epoch 5/10 update 13,100/28,680 loss 2,183 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 66 min epoch 5/10 update 13,200/28,680 loss 1,978 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 41 img/s ETA 62 min epoch 5/10 update 13,300/28,680 loss 1,964 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 41 img/s ETA 63 min epoch 5/10 update 13,400/28,680 loss 2,548 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 35 img/s ETA 73 min epoch 5/10 update 13,500/28,680 loss 2,336 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 36 img/s ETA 70 min epoch 5/10 update 13,600/28,680 loss 2,263 vs chance: fuse 0.48x, sides 0.60x lr 1e-07 40 img/s ETA 63 min epoch 5/10 update 13,700/28,680 loss 2,155 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 62 min epoch 5/10 update 13,800/28,680 loss 2,080 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 42 img/s ETA 59 min epoch 5/10 update 13,900/28,680 loss 2,143 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 41 img/s ETA 60 min epoch 5/10 update 14,000/28,680 loss 2,215 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 38 img/s ETA 65 min epoch 5/10 update 14,100/28,680 loss 2,211 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 38 img/s ETA 63 min epoch 5/10 update 14,200/28,680 loss 1,962 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 39 img/s ETA 63 min epoch 5/10 update 14,300/28,680 loss 2,426 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 38 img/s ETA 64 min --- epoch 5 done, checkpoint saved (63 min this session) epoch 6/10 update 14,400/28,680 loss 2,716 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 32 img/s ETA 75 min epoch 6/10 update 14,500/28,680 loss 2,194 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 38 img/s ETA 61 min epoch 6/10 update 14,600/28,680 loss 2,129 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 38 img/s ETA 62 min epoch 6/10 update 14,700/28,680 loss 2,291 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 36 img/s ETA 65 min epoch 6/10 update 14,800/28,680 loss 2,374 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 37 img/s ETA 63 min epoch 6/10 update 14,900/28,680 loss 2,221 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 58 min epoch 6/10 update 15,000/28,680 loss 2,154 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 56 min epoch 6/10 update 15,100/28,680 loss 1,979 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 38 img/s ETA 59 min epoch 6/10 update 15,200/28,680 loss 2,158 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 39 img/s ETA 58 min epoch 6/10 update 15,300/28,680 loss 1,758 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 45 img/s ETA 50 min epoch 6/10 update 15,400/28,680 loss 1,860 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 44 img/s ETA 50 min epoch 6/10 update 15,500/28,680 loss 1,934 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 44 img/s ETA 50 min epoch 6/10 update 15,600/28,680 loss 2,197 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 39 img/s ETA 56 min epoch 6/10 update 15,700/28,680 loss 2,421 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 38 img/s ETA 56 min epoch 6/10 update 15,800/28,680 loss 2,128 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 44 img/s ETA 49 min epoch 6/10 update 15,900/28,680 loss 1,941 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 43 img/s ETA 49 min epoch 6/10 update 16,000/28,680 loss 1,958 vs chance: fuse 0.48x, sides 0.60x lr 1e-07 45 img/s ETA 47 min epoch 6/10 update 16,100/28,680 loss 1,822 vs chance: fuse 0.48x, sides 0.60x lr 1e-07 43 img/s ETA 49 min epoch 6/10 update 16,200/28,680 loss 2,002 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 43 img/s ETA 49 min epoch 6/10 update 16,300/28,680 loss 2,601 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 34 img/s ETA 61 min epoch 6/10 update 16,400/28,680 loss 2,255 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 38 img/s ETA 55 min epoch 6/10 update 16,500/28,680 loss 2,066 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 40 img/s ETA 51 min epoch 6/10 update 16,600/28,680 loss 2,027 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 43 img/s ETA 46 min epoch 6/10 update 16,700/28,680 loss 2,560 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 34 img/s ETA 58 min epoch 6/10 update 16,800/28,680 loss 2,496 vs chance: fuse 0.46x, sides 0.58x lr 1e-07 34 img/s ETA 58 min epoch 6/10 update 16,900/28,680 loss 2,198 vs chance: fuse 0.46x, sides 0.58x lr 1e-07 38 img/s ETA 52 min epoch 6/10 update 17,000/28,680 loss 2,160 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 38 img/s ETA 51 min epoch 6/10 update 17,100/28,680 loss 2,483 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 35 img/s ETA 55 min epoch 6/10 update 17,200/28,680 loss 2,139 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 48 min --- epoch 6 done, checkpoint saved (76 min this session) epoch 7/10 update 17,300/28,680 loss 2,238 vs chance: fuse 0.46x, sides 0.57x lr 1e-07 34 img/s ETA 56 min epoch 7/10 update 17,400/28,680 loss 2,407 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 38 img/s ETA 50 min epoch 7/10 update 17,500/28,680 loss 2,154 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 41 img/s ETA 46 min epoch 7/10 update 17,600/28,680 loss 2,297 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 36 img/s ETA 51 min epoch 7/10 update 17,700/28,680 loss 2,224 vs chance: fuse 0.48x, sides 0.60x lr 1e-07 39 img/s ETA 47 min epoch 7/10 update 17,800/28,680 loss 2,140 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 38 img/s ETA 48 min epoch 7/10 update 17,900/28,680 loss 2,151 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 39 img/s ETA 46 min epoch 7/10 update 18,000/28,680 loss 2,476 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 37 img/s ETA 49 min epoch 7/10 update 18,100/28,680 loss 2,213 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 40 img/s ETA 45 min epoch 7/10 update 18,200/28,680 loss 1,898 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 45 img/s ETA 39 min epoch 7/10 update 18,300/28,680 loss 2,117 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 40 img/s ETA 43 min epoch 7/10 update 18,400/28,680 loss 1,817 vs chance: fuse 0.46x, sides 0.58x lr 1e-07 41 img/s ETA 42 min epoch 7/10 update 18,500/28,680 loss 2,088 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 40 img/s ETA 43 min epoch 7/10 update 18,600/28,680 loss 2,455 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 35 img/s ETA 48 min epoch 7/10 update 18,700/28,680 loss 1,795 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 46 img/s ETA 36 min epoch 7/10 update 18,800/28,680 loss 2,108 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 42 img/s ETA 39 min epoch 7/10 update 18,900/28,680 loss 2,256 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 37 img/s ETA 44 min epoch 7/10 update 19,000/28,680 loss 2,152 vs chance: fuse 0.49x, sides 0.60x lr 1e-07 41 img/s ETA 40 min epoch 7/10 update 19,100/28,680 loss 2,337 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 33 img/s ETA 48 min epoch 7/10 update 19,200/28,680 loss 2,112 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 40 img/s ETA 39 min epoch 7/10 update 19,300/28,680 loss 2,184 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 39 min epoch 7/10 update 19,400/28,680 loss 1,978 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 41 img/s ETA 38 min epoch 7/10 update 19,500/28,680 loss 2,338 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 39 img/s ETA 39 min epoch 7/10 update 19,600/28,680 loss 2,230 vs chance: fuse 0.48x, sides 0.59x lr 1e-07 40 img/s ETA 38 min epoch 7/10 update 19,700/28,680 loss 2,096 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 39 img/s ETA 38 min epoch 7/10 update 19,800/28,680 loss 1,951 vs chance: fuse 0.47x, sides 0.59x lr 1e-07 45 img/s ETA 33 min epoch 7/10 update 19,900/28,680 loss 2,260 vs chance: fuse 0.47x, sides 0.58x lr 1e-07 37 img/s ETA 40 min epoch 7/10 update 20,000/28,680 loss 2,023 vs chance: fuse 0.46x, sides 0.57x lr 1e-07 39 img/s ETA 37 min --- epoch 7 done, checkpoint saved (88 min this session) epoch 8/10 update 20,100/28,680 loss 1,734 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 45 img/s ETA 32 min epoch 8/10 update 20,200/28,680 loss 2,249 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 37 img/s ETA 38 min epoch 8/10 update 20,300/28,680 loss 2,631 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 37 img/s ETA 38 min epoch 8/10 update 20,400/28,680 loss 2,114 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 40 img/s ETA 34 min epoch 8/10 update 20,500/28,680 loss 2,505 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 37 img/s ETA 37 min epoch 8/10 update 20,600/28,680 loss 2,150 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 42 img/s ETA 32 min epoch 8/10 update 20,700/28,680 loss 2,276 vs chance: fuse 0.46x, sides 0.57x lr 1e-08 37 img/s ETA 36 min epoch 8/10 update 20,800/28,680 loss 1,961 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 41 img/s ETA 32 min epoch 8/10 update 20,900/28,680 loss 1,993 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 39 img/s ETA 33 min epoch 8/10 update 21,000/28,680 loss 2,055 vs chance: fuse 0.45x, sides 0.56x lr 1e-08 39 img/s ETA 33 min epoch 8/10 update 21,100/28,680 loss 2,240 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 39 img/s ETA 33 min epoch 8/10 update 21,200/28,680 loss 2,314 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 36 img/s ETA 35 min epoch 8/10 update 21,300/28,680 loss 2,340 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 40 img/s ETA 31 min epoch 8/10 update 21,400/28,680 loss 2,412 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 35 img/s ETA 34 min epoch 8/10 update 21,500/28,680 loss 2,026 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 41 img/s ETA 29 min epoch 8/10 update 21,600/28,680 loss 2,131 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 38 img/s ETA 31 min epoch 8/10 update 21,700/28,680 loss 2,055 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 40 img/s ETA 29 min epoch 8/10 update 21,800/28,680 loss 2,195 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 38 img/s ETA 30 min epoch 8/10 update 21,900/28,680 loss 2,015 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 40 img/s ETA 28 min epoch 8/10 update 22,000/28,680 loss 1,749 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 49 img/s ETA 23 min epoch 8/10 update 22,100/28,680 loss 1,930 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 39 img/s ETA 28 min epoch 8/10 update 22,200/28,680 loss 2,349 vs chance: fuse 0.46x, sides 0.57x lr 1e-08 37 img/s ETA 29 min epoch 8/10 update 22,300/28,680 loss 1,595 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 48 img/s ETA 22 min epoch 8/10 update 22,400/28,680 loss 1,972 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 40 img/s ETA 26 min epoch 8/10 update 22,500/28,680 loss 2,024 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 39 img/s ETA 27 min epoch 8/10 update 22,600/28,680 loss 2,534 vs chance: fuse 0.46x, sides 0.57x lr 1e-08 34 img/s ETA 30 min epoch 8/10 update 22,700/28,680 loss 1,911 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 43 img/s ETA 23 min epoch 8/10 update 22,800/28,680 loss 2,371 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 35 img/s ETA 28 min epoch 8/10 update 22,900/28,680 loss 2,433 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 35 img/s ETA 27 min --- epoch 8 done, checkpoint saved (100 min this session) epoch 9/10 update 23,000/28,680 loss 2,054 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 44 img/s ETA 21 min epoch 9/10 update 23,100/28,680 loss 2,230 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 37 img/s ETA 25 min epoch 9/10 update 23,200/28,680 loss 2,426 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 38 img/s ETA 24 min epoch 9/10 update 23,300/28,680 loss 2,111 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 39 img/s ETA 23 min epoch 9/10 update 23,400/28,680 loss 2,000 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 43 img/s ETA 20 min epoch 9/10 update 23,500/28,680 loss 2,231 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 37 img/s ETA 23 min epoch 9/10 update 23,600/28,680 loss 2,542 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 34 img/s ETA 25 min epoch 9/10 update 23,700/28,680 loss 2,266 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 36 img/s ETA 23 min epoch 9/10 update 23,800/28,680 loss 2,082 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 42 img/s ETA 19 min epoch 9/10 update 23,900/28,680 loss 2,092 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 42 img/s ETA 19 min epoch 9/10 update 24,000/28,680 loss 1,869 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 44 img/s ETA 18 min epoch 9/10 update 24,100/28,680 loss 1,995 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 42 img/s ETA 18 min epoch 9/10 update 24,200/28,680 loss 2,096 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 37 img/s ETA 20 min epoch 9/10 update 24,300/28,680 loss 1,890 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 44 img/s ETA 17 min epoch 9/10 update 24,400/28,680 loss 2,412 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 35 img/s ETA 20 min epoch 9/10 update 24,500/28,680 loss 2,020 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 40 img/s ETA 18 min epoch 9/10 update 24,600/28,680 loss 2,155 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 38 img/s ETA 18 min epoch 9/10 update 24,700/28,680 loss 2,419 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 36 img/s ETA 19 min epoch 9/10 update 24,800/28,680 loss 2,066 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 39 img/s ETA 17 min epoch 9/10 update 24,900/28,680 loss 2,217 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 38 img/s ETA 17 min epoch 9/10 update 25,000/28,680 loss 2,015 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 45 img/s ETA 14 min epoch 9/10 update 25,100/28,680 loss 2,261 vs chance: fuse 0.46x, sides 0.57x lr 1e-08 34 img/s ETA 17 min epoch 9/10 update 25,200/28,680 loss 2,450 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 37 img/s ETA 16 min epoch 9/10 update 25,300/28,680 loss 1,845 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 43 img/s ETA 13 min epoch 9/10 update 25,400/28,680 loss 2,142 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 39 img/s ETA 14 min epoch 9/10 update 25,500/28,680 loss 2,051 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 40 img/s ETA 13 min epoch 9/10 update 25,600/28,680 loss 2,311 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 41 img/s ETA 12 min epoch 9/10 update 25,700/28,680 loss 2,127 vs chance: fuse 0.46x, sides 0.57x lr 1e-08 39 img/s ETA 13 min epoch 9/10 update 25,800/28,680 loss 2,000 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 40 img/s ETA 12 min --- epoch 9 done, checkpoint saved (112 min this session) epoch 10/10 update 25,900/28,680 loss 2,306 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 38 img/s ETA 12 min epoch 10/10 update 26,000/28,680 loss 1,855 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 42 img/s ETA 11 min epoch 10/10 update 26,100/28,680 loss 2,404 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 37 img/s ETA 12 min epoch 10/10 update 26,200/28,680 loss 2,108 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 40 img/s ETA 10 min epoch 10/10 update 26,300/28,680 loss 2,279 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 39 img/s ETA 10 min epoch 10/10 update 26,400/28,680 loss 2,056 vs chance: fuse 0.45x, sides 0.57x lr 1e-08 38 img/s ETA 10 min epoch 10/10 update 26,500/28,680 loss 2,119 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 38 img/s ETA 10 min epoch 10/10 update 26,600/28,680 loss 2,405 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 37 img/s ETA 9 min epoch 10/10 update 26,700/28,680 loss 2,206 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 35 img/s ETA 9 min epoch 10/10 update 26,800/28,680 loss 1,873 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 42 img/s ETA 8 min epoch 10/10 update 26,900/28,680 loss 2,489 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 39 img/s ETA 8 min epoch 10/10 update 27,000/28,680 loss 1,875 vs chance: fuse 0.48x, sides 0.60x lr 1e-08 42 img/s ETA 7 min epoch 10/10 update 27,100/28,680 loss 2,143 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 43 img/s ETA 6 min epoch 10/10 update 27,200/28,680 loss 1,769 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 43 img/s ETA 6 min epoch 10/10 update 27,300/28,680 loss 2,373 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 36 img/s ETA 6 min epoch 10/10 update 27,400/28,680 loss 2,210 vs chance: fuse 0.48x, sides 0.59x lr 1e-08 40 img/s ETA 5 min epoch 10/10 update 27,500/28,680 loss 2,227 vs chance: fuse 0.45x, sides 0.57x lr 1e-08 36 img/s ETA 5 min epoch 10/10 update 27,600/28,680 loss 2,160 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 37 img/s ETA 5 min epoch 10/10 update 27,700/28,680 loss 2,131 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 41 img/s ETA 4 min epoch 10/10 update 27,800/28,680 loss 1,840 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 43 img/s ETA 3 min epoch 10/10 update 27,900/28,680 loss 2,285 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 38 img/s ETA 3 min epoch 10/10 update 28,000/28,680 loss 2,079 vs chance: fuse 0.47x, sides 0.58x lr 1e-08 38 img/s ETA 3 min epoch 10/10 update 28,100/28,680 loss 2,438 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 38 img/s ETA 3 min epoch 10/10 update 28,200/28,680 loss 2,022 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 41 img/s ETA 2 min epoch 10/10 update 28,300/28,680 loss 2,259 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 39 img/s ETA 2 min epoch 10/10 update 28,400/28,680 loss 2,030 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 41 img/s ETA 1 min epoch 10/10 update 28,500/28,680 loss 2,123 vs chance: fuse 0.46x, sides 0.58x lr 1e-08 38 img/s ETA 1 min epoch 10/10 update 28,600/28,680 loss 2,281 vs chance: fuse 0.47x, sides 0.59x lr 1e-08 37 img/s ETA 0 min --- epoch 10 done, checkpoint saved (125 min this session) training complete
import pandas as pd
hist = pd.read_csv(LOG_PATH)
if len(hist):
fig, axes = plt.subplots(1, 3, figsize=(17, 4.5))
window = max(1, len(hist) // 50)
x = hist['update']
axes[0].plot(x, hist['loss'], alpha=0.3, color='C0')
axes[0].plot(x, hist['loss'].rolling(window, min_periods=1).mean(), color='C0')
axes[0].set(title='total loss (sum of 6 outputs)', xlabel='update', yscale='log')
for k in ['dsn1', 'dsn2', 'dsn3', 'dsn4', 'dsn5', 'fuse']:
axes[1].plot(x, (hist[k] / hist['chance']).rolling(window, min_periods=1).mean(), label=k)
axes[1].axhline(1, color='gray', ls='--', lw=1)
axes[1].set(title='loss per output / chance level (1 = guessing)', xlabel='update')
axes[1].legend()
axes[2].plot(x, hist['lr'])
axes[2].set(title='learning rate', xlabel='update', yscale='log')
for ax in axes:
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(FIG_DIR / 'training_curves.png', dpi=100)
plt.show()
7. Export weights for hed-tensor¶
hed_bsds500.weights.h5: drop-in replacement forhed_tensor/weights/hed_bsds500.weights.h5.hed_caffe_layout.npz: the same numbers under the original Caffe layer names (conv1_1/W,score-dsn1/b,new-score-weighting/W, ...). Kernels are Keras HWIO; for PyTorch useW.transpose(3, 2, 0, 1).hed_config.json: preprocessing and training settings.
import json
WEIGHTS_FILE = EXPORT / 'hed_bsds500.weights.h5'
model.save_weights(WEIGHTS_FILE)
CAFFE_CONV_NAMES = ['conv1_1', 'conv1_2', 'conv2_1', 'conv2_2', 'conv3_1', 'conv3_2', 'conv3_3',
'conv4_1', 'conv4_2', 'conv4_3', 'conv5_1', 'conv5_2', 'conv5_3']
named = list(zip(CAFFE_CONV_NAMES, [l for blk in model.blocks() for l in blk.layers if isinstance(l, layers.Conv2D)]))
named += [(f'score-dsn{k}', dsn) for k, dsn in enumerate(model.dsns(), 1)] + [('new-score-weighting', model.fuse)]
np.savez(EXPORT / 'hed_caffe_layout.npz', **{f'{n}/{p}': getattr(l, a).numpy()
for n, l in named for p, a in (('W', 'kernel'), ('b', 'bias'))})
hist = pd.read_csv(LOG_PATH)
CONFIG = {
'model': 'HED (Xie & Tu, ICCV 2015), VGG16 backbone, TensorFlow/Keras',
'weights_file': WEIGHTS_FILE.name,
'compatible_with': 'hed-tensor (hed_tensor.HED: block1-5, dsn1-5, fuse)',
'input': {'layout': 'NHWC', 'color_order': 'BGR', 'mean_bgr': MEAN_BGR.tolist(),
'pixel_range': '0-255 minus mean, no scaling', 'size': 'any H x W'},
'output': 'sigmoid(fuse), H x W x 1 in [0, 1]; call(training=True) gives the 6 logits',
'npz_layout': 'Keras HWIO kernels; PyTorch Conv2d weight = W.transpose(3, 2, 0, 1)',
'training': {
'complete': TRAINING_COMPLETE, 'epochs_done': int(EPOCH.numpy()), 'updates_done': int(UPDATE.numpy()),
'images_per_update': BATCH_SIZE, 'base_lr': BASE_LR, 'lr_step_updates': LR_STEP_UPDATES, 'lr_gamma': LR_GAMMA,
'momentum': MOMENTUM, 'weight_decay': WEIGHT_DECAY, 'label_threshold': LABEL_THRESHOLD,
'lr_mult': {'conv1-4': 1, 'conv5': 100, 'dsn': 0.01, 'fuse': 0.001, 'bias': '2x weight'},
'dataset': f'HED-BSDS augmented train set, {len(pairs):,} pairs',
'final_fuse_loss_vs_chance': round(float(hist['fuse'].iloc[-1] / hist['chance'].iloc[-1]), 4) if len(hist) else None,
'smoke_test': SMOKE_TEST,
},
'environment': {'tensorflow': tf.__version__, 'keras': keras.__version__, 'gpus': GPU_NAMES,
'date': time.strftime('%Y-%m-%d')},
}
(EXPORT / 'hed_config.json').write_text(json.dumps(CONFIG, indent=2))
for p in sorted(EXPORT.iterdir()):
print(f'{p.stat().st_size / 1e6:7.1f} MB {p}')
58.9 MB /kaggle/working/export/hed_bsds500.weights.h5
58.9 MB /kaggle/working/export/hed_caffe_layout.npz
0.0 MB /kaggle/working/export/hed_config.json
Check: the exported file loads in hed-tensor and gives the same output¶
The library prints a warning instead of failing when weights don't load, so this compares actual outputs.
lib_model = hed_tensor.HED(pretrained=False, weights=str(WEIGHTS_FILE))
_, x = load_test(TEST_PATHS[0])
diff = float(np.abs(lib_model(x, training=False).numpy() - model(x, training=False).numpy()).max())
print(f'max difference library vs notebook model: {diff:.2e}')
assert diff < 1e-5, 'exported weights did not load correctly into hed_tensor.HED'
del lib_model
Loaded weights from: /kaggle/working/export/hed_bsds500.weights.h5 max difference library vs notebook model: 0.00e+00
8. Predict the 200 BSDS500 test images¶
Fused output and the 5 side outputs as PNGs (BSDS benchmark format), zipped in predictions/. Comparison with other models is done separately.
def save_map(prob, path):
path.parent.mkdir(parents=True, exist_ok=True)
Image.fromarray(np.round(np.clip(prob, 0, 1) * 255).astype(np.uint8)).save(path)
t0 = time.time()
for path in TEST_PATHS:
maps = [m.numpy()[0, ..., 0] for m in predict_all(load_test(path)[1])]
save_map(maps[-1], PRED / 'fuse' / f'{path.stem}.png')
for k, m in enumerate(maps[:5], 1):
save_map(m, PRED / f'dsn{k}' / f'{path.stem}.png')
print(f'{len(TEST_PATHS)} test images in {time.time() - t0:.0f}s')
(WORK / 'predictions').mkdir(exist_ok=True)
shutil.make_archive(str(WORK / 'predictions' / 'ours_fuse'), 'zip', PRED / 'fuse')
shutil.make_archive(str(WORK / 'predictions' / 'ours_all_outputs'), 'zip', PRED)
fig, axes = plt.subplots(3, 7, figsize=(21, 7.5))
for r, path in enumerate(TEST_PATHS[:3]):
panels = [load_test(path)[0]] + [np.array(Image.open(PRED / n / f'{path.stem}.png'))
for n in ['dsn1', 'dsn2', 'dsn3', 'dsn4', 'dsn5', 'fuse']]
for c, (p, title) in enumerate(zip(panels, ['image', 'dsn1', 'dsn2', 'dsn3', 'dsn4', 'dsn5', 'fuse'])):
axes[r, c].imshow(p if c == 0 else 255 - p, cmap=None if c == 0 else 'gray', vmin=0, vmax=255)
axes[r, c].axis('off')
if r == 0:
axes[r, c].set_title(title)
fig.tight_layout()
fig.savefig(FIG_DIR / 'side_outputs.png', dpi=90)
plt.show()
200 test images in 28s
9. Output summary¶
Download export/ for the library. To update hed-tensor: replace hed_tensor/weights/hed_bsds500.weights.h5 with the exported file and bump the version.
files = sorted(p for p in WORK.rglob('*') if p.is_file())
print(f'{len(files)} files, {sum(p.stat().st_size for p in files) / 1e9:.2f} GB in {WORK}')
for p in files:
if p.parent.name != 'checkpoints':
print(f'{p.stat().st_size / 1e6:8.1f} MB {p.relative_to(WORK)}')
24 files, 0.46 GB in /kaggle/working
1.9 MB __notebook__.ipynb
58.9 MB export/hed_bsds500.weights.h5
58.9 MB export/hed_caffe_layout.npz
0.0 MB export/hed_config.json
0.5 MB figures/progress_epoch_01.png
0.5 MB figures/progress_epoch_02.png
0.5 MB figures/progress_epoch_03.png
0.5 MB figures/progress_epoch_04.png
0.5 MB figures/progress_epoch_05.png
0.5 MB figures/progress_epoch_06.png
0.5 MB figures/progress_epoch_07.png
0.5 MB figures/progress_epoch_08.png
0.5 MB figures/progress_epoch_09.png
0.5 MB figures/progress_epoch_10.png
1.1 MB figures/side_outputs.png
0.1 MB figures/training_curves.png
0.0 MB logs/train_log.csv
87.1 MB predictions/ours_all_outputs.zip
14.7 MB predictions/ours_fuse.zip