cleanups and adjust learning rate for fp16

examples/mlperf/losses.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from examples.mlperf.metrics import dice_score
 from tinygrad import Tensor
 

examples/mlperf/model_train.py

@@ -351,7 +351,7 @@ def train_retinanet():
   from extra.lr_scheduler import LambdaLR
   from pycocotools.coco import COCO
   from pycocotools.cocoeval import COCOeval
-  from tinygrad.helpers import colored, Context, DEBUG
+  from tinygrad.helpers import colored, Context
   from tinygrad.nn.optim import Optimizer
   from typing import Iterator
   import extra.models.retinanet as retinanet
@@ -419,7 +419,7 @@ def train_retinanet():
   config["epochs"] = EPOCHS = getenv("EPOCHS", 4)
   config["train_beam"] = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
   config["eval_beam"] = EVAL_BEAM = getenv("EVAL_BEAM", BEAM.value)
-  config["lr"] = lr = getenv("LR", 0.0001 * (BS / 256))
+  config["lr"] = lr = getenv("LR", 0.00001 * (BS / 256))
   config["lr_warmup_epochs"] = lr_warmup_epochs = getenv("LR_WARMUP_EPOCHS", 1)
   config["lr_warmup_factor"] = lr_warmup_factor = getenv("LR_WARMUP_FACTOR", 1e-3)
   config["loss_scaler"] = loss_scaler = getenv("LOSS_SCALER", 256.0 if dtypes.default_float == dtypes.float16 else 1.0)

extra/models/retinanet.py

@@ -143,14 +143,14 @@ class ClassificationHead:
 
     if Tensor.training:
       assert labels is not None and matches is not None, "labels and matches should be passed in when training"
-      return self._compute_loss(out, labels, matches)
+      return self._compute_loss(out.cast(dtypes.float32), labels, matches)
 
     return out.sigmoid()
 
   def _compute_loss(self, x:Tensor, labels:Tensor, matches:Tensor) -> Tensor:
     labels = ((labels + 1) * (fg_idxs := matches >= 0) - 1).one_hot(num_classes=x.shape[-1])
     valid_idxs = (matches != -2).reshape(matches.shape[0], -1, 1)
-    loss = valid_idxs.where(sigmoid_focal_loss(x.cast(dtypes.float32), labels), 0).sum(-1).sum(-1)
+    loss = valid_idxs.where(sigmoid_focal_loss(x, labels), 0).sum(-1).sum(-1)
     loss = (loss / fg_idxs.sum(-1)).sum() / matches.shape[0]
     return loss
 
@@ -175,7 +175,7 @@ class RegressionHead:
 
   def _compute_loss(self, x:Tensor, bboxes:Tensor, matches:Tensor, anchors:Tensor) -> Tensor:
     mask = (fg_idxs := matches >= 0).reshape(matches.shape[0], -1, 1)
-    x = x.cast(dtypes.float32) * mask
+    x = x * mask
     tgt = self.box_coder.encode(bboxes, anchors) * mask
     loss = l1_loss(x, tgt).sum(-1).sum(-1)
     loss = (loss / fg_idxs.sum(-1)).sum() / matches.shape[0]