More robust tensor conversion and updates in training script

experiments/unet-segmentation/dsb/train_boundaries.py

@@ -9,11 +9,16 @@ def train_boundaries(args):
 
  patch_shape = (1, 256, 256)
  train_loader = get_dsb_loader(
- args.input, patch_shape, split="train",
+ args.input, patch_shape=patch_shape, split="train",
  download=True, boundaries=True, batch_size=args.batch_size
  )
+
+ # Uncomment this for checking the loader.
+ # from torch_em.util.debug import check_loader
+ # check_loader(train_loader, 4)
+
  val_loader = get_dsb_loader(
- args.input, patch_shape, split="test",
+ args.input, patch_shape=patch_shape, split="test",
  boundaries=True, batch_size=args.batch_size
  )
  loss = torch_em.loss.DiceLoss()

torch_em/util/debug.py

@@ -106,6 +106,7 @@ def _check_napari(loader, n_samples, instance_labels, model=None, device=None, r
  v.add_image(y)
  if pred is not None:
  v.add_image(pred)
+
  napari.run()
 
 

torch_em/util/util.py

@@ -68,7 +68,14 @@ def ensure_tensor(tensor, dtype=None):
  if isinstance(tensor, np.ndarray):
  if np.dtype(tensor.dtype) in DTYPE_MAP:
  tensor = tensor.astype(DTYPE_MAP[tensor.dtype])
- tensor = torch.from_numpy(tensor)
+ # Try to convert the tensor, even if it has wrong byte-order
+ try:
+ tensor = torch.from_numpy(tensor)
+ except ValueError:
+ tensor = tensor.view(tensor.dtype.newbyteorder())
+ if np.dtype(tensor.dtype) in DTYPE_MAP:
+ tensor = tensor.astype(DTYPE_MAP[tensor.dtype])
+ tensor = torch.from_numpy(tensor)
 
  assert torch.is_tensor(tensor), f"Cannot convert {type(tensor)} to torch"
  if dtype is not None: