Comfy-Org / Comfy-Org/ComfyUI

Smoothing and average for plot loss graph

Open
#9,748 1 comment 0 reactions 0 assignees View on GitHub
Feature
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 6h
Merged PRs (30d)
155

Description

### Feature Idea

The plot loss graph node as is doesn't really help with the LoRA training process, it's very difficult to see progress on shorter runs (<1000 steps). I'd like to see a smoothed overlay on the loss graph image to better show what's going on without the noise, an average number output would also be helpful so you can determine progress for shorter <500 step runs at a glance

### Existing Solutions

I had chatGPT help with some code but this was minimal changes it came up with and it appears to work for visual aid, i'm not an expert and python is not my language of choice so it may be incorrect but here's what ChatGPT gave me for the plot_loss definition and it's been working well so far
`def plot_loss(self, loss, filename_prefix, prompt=None, extra_pnginfo=None):
loss_values = loss["loss"]
width, height = 800, 480
margin = 40

img = Image.new(
"RGB", (width + margin, height + margin), "white"
) # Extend canvas
draw = ImageDraw.Draw(img)

min_loss, max_loss = min(loss_values), max(loss_values)
scaled_loss = [(l - min_loss) / (max_loss - min_loss) for l in loss_values]

# --- tiny EMA smoothing (key tweak) ---
beta = 0.96 # 0.98–0.995 = smoother
m = scaled_loss[0]
smoothed = []
for v in scaled_loss:
m = beta * m + (1.0 - beta) * v
smoothed.append(m)
# --------------------------------------

steps = len(loss_values)
x_den = max(steps - 1, 1)

# Raw (light gray)
prev = (margin, height - int(scaled_loss[0] * height))
for i, v in enumerate(scaled_loss[1:], start=1):
x = margin + int(i / steps * width)
y = height - int(v * height)
draw.line([prev, (x, y)], fill=(200, 200, 200), width=1)
prev = (x, y)

# Smoothed (blue)
prev = (margin, height - int(smoothed[0] * height))
for i, v in enumerate(smoothed[1:], start=1):
x = margin + int(i / x_den * width)
y = height - int(v * height)
draw.line([prev, (x, y)], fill="blue", width=2)
prev = (x, y)

draw.line([(margin, 0), (margin, height)], fill="black", width=2) # Y-axis
draw.line(
[(margin, height), (width + margin, height)], fill="black", width=2
) # X-axis

font = None
try:
font = ImageFont.truetype("arial.ttf", 12)
except IOError:
font = ImageFont.load_default()

# Add axis labels
draw.text((5, height // 2), "Loss", font=font, fill="black")
draw.text((width // 2, height + 10), "Steps", font=font, fill="black")

# Add min/max loss values
draw.text((margin - 30, 0), f"{max_loss:.2f}", font=font, fill="black")
draw.text(
(margin - 30, height - 10), f"{min_loss:.2f}", font=font, fill="black"
)

# --- Average smoothed value marker on Y-axis ---
avg_scaled = sum(scaled_loss) / max(1, len(scaled_loss)) # 0..1 (same scaling)
avg_loss = min_loss + avg_scaled * (max_loss - min_loss) # original units
y_avg = height - int(avg_scaled * height)

# short tick centered on the Y-axis
draw.line([(margin-6, y_avg), (margin+6, y_avg)], fill="blue", width=2)
# label the value
draw.text((5, y_avg - 7), f"{avg_loss:.3f}", font=font, fill="blue")

# (optional) faint dashed guideline across the plot
# for x in range(margin+8, width+margin, 8):
# draw.line([(x, y_avg), (x+4, y_avg)], fill=(150, 200, 255), width=1)
# -----------------------------------------------

metadata = None
if not args.disable_metadata:
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))

date = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
img.save(
os.path.join(self.output_dir, f"{filename_prefix}_{date}.png"),
pnginfo=metadata,
)
return {
"ui": {
"images": [
{
"filename": f"{filename_prefix}_{date}.png",
"subfolder": "",
"type": "temp",
}
]
}
}`

### Other

Feel free to use the code or snippets of the code provided I relinquish any ownership of it, Will require testing though

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.