File size: 4,752 Bytes
9d33171 | 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 | import datasets
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torch import nn
import torchvision
from tqdm import tqdm
from dataset import EuroSATDataset
import torch.nn.functional as F
# Constants retrieved from:
# https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html
RESNET_50_WEIGHT_MEAN = [0.485, 0.456, 0.406]
RESNET_50_WEIGHT_STD = [0.229, 0.224, 0.225]
DATASET_CFG = {
"rgb": {"hf_id": "blanchon/EuroSAT_RGB", "in_channels": 3},
"msi": {"hf_id": "blanchon/EuroSAT_MSI", "in_channels": 13},
}
def to_chw_tensor(image):
hwc = np.array(image, dtype=np.float32) # HWC typical shape: 64x64x3
chw = torch.from_numpy(hwc).permute(2, 0, 1) # CHW typical shape: 3x64x64
return chw
def build_rgb_transform(train: bool):
ops = [transforms.Resize((224, 224))]
if train:
ops.append(transforms.RandomHorizontalFlip())
ops.extend(
[
transforms.ToTensor(),
transforms.Normalize(RESNET_50_WEIGHT_MEAN, RESNET_50_WEIGHT_STD),
]
)
return transforms.Compose(ops)
def build_msi_transform(train: bool):
def _tf(image):
chw = to_chw_tensor(image)
chw = chw / 10000.0
if train and torch.rand(1).item() < 0.5:
chw = torch.flip(chw, dims=[2])
chw = F.interpolate(
chw.unsqueeze(0), size=(224, 224), mode="bilinear", align_corners=False
).squeeze(0)
return chw
return _tf
def build_dataloaders(
modality: str,
batch_size: int,
num_workers: int,
):
cfg = DATASET_CFG[modality]
ds = datasets.load_dataset(cfg["hf_id"])
in_channels = cfg["in_channels"]
num_classes = ds["train"].features["label"].num_classes
if modality == "rgb":
train_tf = build_rgb_transform(train=True)
eval_tf = build_rgb_transform(train=False)
else:
train_tf = build_msi_transform(train=True)
eval_tf = build_msi_transform(train=False)
train_ds = EuroSATDataset(ds["train"], train_tf)
val_ds = EuroSATDataset(ds["validation"], eval_tf)
train_loader = DataLoader(
train_ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=torch.cuda.is_available(),
)
val_loader = DataLoader(
val_ds,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=torch.cuda.is_available(),
)
return train_loader, val_loader, num_classes, in_channels
# Helper function to get the device CPU or GPU available to train the models.
def get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def build_model(num_classes: int, device: torch.device, in_channels: int) -> nn.Module:
model = torchvision.models.resnet50(weights=None)
if in_channels != 3:
model.conv1 = nn.Conv2d(
in_channels=in_channels,
out_channels=model.conv1.out_channels,
kernel_size=model.conv1.kernel_size,
stride=model.conv1.stride,
padding=model.conv1.padding,
bias=False,
)
model.fc = nn.Linear(model.fc.in_features, num_classes)
return model.to(device)
def train_one_epoch(
model: nn.Module,
loader: DataLoader,
criterion: nn.Module,
optimizer: torch.optim.Optimizer,
device: torch.device,
):
model.train()
total_loss = 0.0
n = 0
for images, labels in tqdm(loader, desc="train", leave=False):
images = images.to(device)
labels = labels.to(device, dtype=torch.long)
optimizer.zero_grad()
logits = model(images)
loss = criterion(logits, labels)
loss.backward()
optimizer.step()
batch_n = labels.size(0)
total_loss += loss.item() * batch_n
n += batch_n
train_loss = total_loss / max(n, 1) # max(n, 1) to avoid division by zero
return train_loss
@torch.no_grad()
def evaluate(
model: nn.Module, loader: DataLoader, criterion: nn.Module, device: torch.device
):
model.eval()
total_loss, correct, total = 0.0, 0, 0
for images, labels in loader:
images = images.to(device)
labels = labels.to(device)
logits = model(images)
loss = criterion(logits, labels)
total_loss += loss.item() * labels.size(0)
correct += (logits.argmax(1) == labels).sum().item()
total += labels.size(0)
val_loss = total_loss / total
val_acc = correct / total
return val_loss, val_acc
|