How to specify loss function?

[rosenfeldamir]

The trainer module accepts a string for a loss function. Can I pass a function instead?
Otherwise, where can I find the mapping between these strings and actual pytorch loss functions?
For example, how do I use the CrossEntropyLoss?

[ncullen93]

Here is the function:

def _validate_loss_input(loss):
    dir_f = dir(F)
    loss_fns = [d.lower() for d in dir_f]
    if isinstance(loss, str):
        if loss.lower() == 'unconstrained':
            return lambda x: x
        elif loss.lower() == 'unconstrained_sum':
            return lambda x: x.sum()
        elif loss.lower() == 'unconstrained_mean':
            return lambda x: x.mean()
        else:
            try:
                str_idx = loss_fns.index(loss.lower())
            except:
                raise ValueError('Invalid loss string input - must match pytorch function.')
            return getattr(F, dir(F)[str_idx])
    elif callable(loss):
        return loss
    else:
        raise ValueError('Invalid loss input')

You can see it matches the function names in torch.nn.functional (doesn't matter if the given case is wrong). You can also just pass in a callable function that takes the same arguments as something like F.nll_loss and it will use that.

[rosenfeldamir]

Thanks!