ImagesLoss

Hi,

I’m very excited about this work! I was wondering if ImagesLoss with Sinkhorn divergence would be usable (when it’s ready) as a reconstruction loss for an image autoencoder for representation learning?

Pseudocode to code to illustrate what I was thinking about:

encoder = ConvNet()
decoder = DeconvNet()
optimizer = Optimizer(chain(encoder.parameters(), decoder.parameters()))
loss_fn = ImagesLoss(loss="sinkhorn", p=2, blur=.05)

for images in dataset:
  optimizer.zero_grad()
  encoding = encoder(images)
  image_preds = decoder(encoding)
  loss = loss_fn(images, image_preds)
  loss.backward()
  optimizer.step()

Thanks, Agost

[jeanfeydy]

Hi @abiro ,

Sure! However, please keep in mind that OT theory is always concerned with measures, distributions of mass. In your example, loss_fn(images, images_preds) would return a vector of Wasserstein-like distances between each pair of original and encoded-decoded image, all of them being understood as density maps that represent measures in the rectangular image domain. (As a side note, before doing a .backward(), you would have to add a loss = loss.sum().) Consequently, ImagesLoss and VolumesLoss will be best suited to the processing of segmentation masks in medical imaging and computer vision.

Note that using OT on raw image data is asking for trouble: measure theory handles “black” and “white” as “absence” and “presence” of mass, which is rarely what researchers want to do with natural images. Fortunately, as discussed for instance in the color transfer tutorial, you can perfectly represent an image as a point cloud in some arbitrary feature space, and use SamplesLoss to define a relevant processing. In your example, a simple way of defining a transport-based loss functions between two “natural” images is to represent them as point clouds (x, y, r, g, b) in R^5 and feed them to SamplesLoss. Of course, you could then tune the ratio between the “geometric features” (x, y) and the color coordinates (r, g, b) or use a perceptual color space to improve your results.

Best regards,

Jean

[JunMa11]

Hi @jeanfeydy ,

Thanks for your quick and detailed reply very much.

I’m confused when comparing the above to_pointcloud function and RGB_cloud function in the color transfer demo.

def RGB_cloud(fname, sampling=1, dtype=torch.FloatTensor) :
    A = load_image(fname)
    A = A[::sampling, ::sampling, :]
    return torch.from_numpy(A).type(dtype).view(-1,3)

If inputting a commonly used RGB 2D/CT 3D image with shape=(256,256,3) or (512,512,100), respectively, the expected point cloud outputs are (x, y, r, g, b), (x, y, z, CT_value), right?

However, for the function RGB_cloud, the output shape is (256256, 3), while the output shape of to_pointcloud is (256256, 2). I’m confused about the differences.

Specifically, I want to use the Wasserstein distance in the following two settings.

1. compute the Wasserstein distance between the label and softmax outputs

trainloader = DataLoader(data_train, batch_size, shuffle=True,  num_workers=4)
net = UNet(n_channels=1, n_classes=num_classes, normalization='batchnorm', has_dropout=True)
net = net.cuda()
net.train()

loss_fn = SamplesLoss("sinkhorn", p=2, blur=0.01)
for epoch_num in range(max_epoch):
    for i_batch, sampled_batch in enumerate(trainloader):
        volume_batch, label_batch = sampled_batch['image'], sampled_batch['label']
        # outputs.shape = (b, c, x, y, z), label_shape = (b, 1, x, y, z)
        outputs = net(volume_batch)
        outputs_soft = F.softmax(outputs, dim=1)
        # how should I pass the outputs_soft and label_batch into the loss_fn?
        loss = loss_fn(outputs_soft ?, label_batch?)

Question: how should I pass the outputs_soft and label_batch into the loss_fn?

outputs_soft.shape = (b, c, x, y, z); label_distance_map.shape= (b, 1, x, y, z); or (b, c, x, y, z)- one-hot encoding

2. compute the Wasserstein distance between two distance maps

L1/L2 norm and KL-divergence is commonly used in this setting. Now, I want to try Wasserstein distance.

loss_fn = SamplesLoss("sinkhorn", p=2, blur=0.01)
for epoch_num in range(max_epoch):
    for i_batch, sampled_batch in enumerate(trainloader):
        volume_batch, label_distance_map = sampled_batch['image'], sampled_batch['label']
        # outputs.shape = (b, c, x, y, z), label_shape = (b, c, x, y, z)
        # how should I pass the outputs_soft and label_batch into the loss_fn?
        loss = loss_fn(outputs ?, label_distance_map?)

Question: how should I pass the outputs and label_distance_map into the loss_fn?

outputs.shape = label_distance_map.shape= (b, c, x, y, z);

Note: here the outputs is the logits rather than the softmax outputs. label_distance_map is distance transform of the ground truth by scipy.

Looking forward to your reply at your convenience.

Best regards, Jun