Spaces:
Running
Running
File size: 6,411 Bytes
3b49518 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
# Copyright (c) EPFL VILAB.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# Based on timm, DeiT, DINO, MoCo-v3, BEiT, MAE-priv and MAE code bases
# https://github.com/rwightman/pytorch-image-models/tree/master/timm
# https://github.com/facebookresearch/deit
# https://github.com/facebookresearch/dino
# https://github.com/facebookresearch/moco-v3
# https://github.com/microsoft/unilm/tree/master/beit
# https://github.com/BUPT-PRIV/MAE-priv
# https://github.com/facebookresearch/mae
# --------------------------------------------------------
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
class MaskedCrossEntropyLoss(nn.Module):
"""Cross-entropy loss with masking
:param patch_size: Patch size
:param stride: Stride of task / modality
:param label_smoothing: Amount of smoothing in the loss (default is 0.0)
"""
def __init__(self, patch_size: int = 16, stride: int = 1, label_smoothing : float = 0.0):
super().__init__()
self.patch_size = patch_size
self.stride = stride
self.scale_factor = patch_size // stride
self.label_smoothing = label_smoothing
def forward(self, input, target, mask=None):
loss = F.cross_entropy(input, target, reduction='none', label_smoothing=self.label_smoothing)
if mask is not None:
if mask.sum() == 0:
return torch.tensor(0).to(loss.device)
H, W = input.shape[-2:]
nh, nw = H // self.scale_factor, W // self.scale_factor
# Resize mask and upsample
mask = rearrange(mask, "b (nh nw) -> b nh nw", nh=nh, nw=nw)
mask = F.interpolate(mask.unsqueeze(1).float(), size=(H, W), mode='nearest').squeeze(1)
loss = loss * mask
# Compute mean per sample
loss = loss.flatten(start_dim=1).sum(dim=1) / mask.flatten(start_dim=1).sum(dim=1)
loss = loss.nanmean() # Account for zero masks
else:
loss = loss.mean() # If this is ever nan, we want it to stop training
return loss
class MaskedMSELoss(nn.Module):
"""L1 loss with masking
:param patch_size: Patch size
:param stride: Stride of task / modality
:param norm_pix: Normalized pixel loss
"""
def __init__(self, patch_size: int = 16, stride: int = 1, norm_pix=False):
super().__init__()
self.patch_size = patch_size
self.stride = stride
self.scale_factor = patch_size // stride
self.norm_pix = norm_pix
def patchify(self, imgs, nh, nw):
p = self.scale_factor
x = rearrange(imgs, "b c (nh p1) (nw p2) -> b (nh nw) (p1 p2 c)", nh=nh, nw=nw, p1=p, p2=p)
return x
def unpatchify(self, x, nh, nw):
p = self.scale_factor
imgs = rearrange(x, "b (nh nw) (p1 p2 c) -> b c (nh p1) (nw p2)", nh=nh, nw=nw, p1=p, p2=p)
return imgs
def forward(self, input, target, mask=None):
H, W = input.shape[-2:]
nh, nw = H // self.scale_factor, W // self.scale_factor
if self.norm_pix:
target = self.patchify(target, nh, nw)
mean = target.mean(dim=-1, keepdim=True)
var = target.var(dim=-1, keepdim=True)
eps = 1e-6
target = (target - mean) / torch.sqrt(var + eps)
target = self.unpatchify(target, nh, nw)
loss = F.mse_loss(input, target, reduction='none')
if mask is not None:
if mask.sum() == 0:
return torch.tensor(0).to(loss.device)
# Resize mask and upsample
mask = rearrange(mask, "b (nh nw) -> b nh nw", nh=nh, nw=nw)
mask = F.interpolate(mask.unsqueeze(1).float(), size=(H, W), mode='nearest').squeeze(1)
loss = loss.mean(dim=1) # B, C, H, W -> B, H, W
loss = loss * mask
# Compute mean per sample
loss = loss.flatten(start_dim=1).sum(dim=1) / mask.flatten(start_dim=1).sum(dim=1)
loss = loss.nanmean() # Account for zero masks
else:
loss = loss.mean() # If this is ever nan, we want it to stop training
return loss
class MaskedL1Loss(nn.Module):
"""L1 loss with masking
:param patch_size: Patch size
:param stride: Stride of task / modality
:param norm_pix: Normalized pixel loss
"""
def __init__(self, patch_size: int = 16, stride: int = 1, norm_pix=False):
super().__init__()
self.patch_size = patch_size
self.stride = stride
self.scale_factor = patch_size // stride
self.norm_pix = norm_pix
def patchify(self, imgs, nh, nw):
p = self.scale_factor
x = rearrange(imgs, "b c (nh p1) (nw p2) -> b (nh nw) (p1 p2 c)", nh=nh, nw=nw, p1=p, p2=p)
return x
def unpatchify(self, x, nh, nw):
p = self.scale_factor
imgs = rearrange(x, "b (nh nw) (p1 p2 c) -> b c (nh p1) (nw p2)", nh=nh, nw=nw, p1=p, p2=p)
return imgs
def forward(self, input, target, mask=None):
H, W = input.shape[-2:]
nh, nw = H // self.scale_factor, W // self.scale_factor
if self.norm_pix:
target = self.patchify(target, nh, nw)
mean = target.mean(dim=-1, keepdim=True)
var = target.var(dim=-1, keepdim=True)
eps = 1e-6
target = (target - mean) / torch.sqrt(var + eps)
target = self.unpatchify(target, nh, nw)
loss = F.l1_loss(input, target, reduction='none')
if mask is not None:
if mask.sum() == 0:
return torch.tensor(0).to(loss.device)
# Resize mask and upsample
mask = rearrange(mask, "b (nh nw) -> b nh nw", nh=nh, nw=nw)
mask = F.interpolate(mask.unsqueeze(1).float(), size=(H, W), mode='nearest').squeeze(1)
loss = loss.mean(dim=1) # B, C, H, W -> B, H, W
loss = loss * mask
# Compute mean per sample
loss = loss.flatten(start_dim=1).sum(dim=1) / mask.flatten(start_dim=1).sum(dim=1)
loss = loss.nanmean() # Account for zero masks
else:
loss = loss.mean() # If this is ever nan, we want it to stop training
return loss
|