-
Notifications
You must be signed in to change notification settings - Fork 4
/
models.py
61 lines (51 loc) · 1.76 KB
/
models.py
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
import torch
import torch.nn as nn
import torchvision
import lightning.pytorch as pl
from metrics import SMAPIoUMetric
class SegModel(pl.LightningModule):
def __init__(self):
super(SegModel, self).__init__()
self.learning_rate = 1e-3
self.net = torchvision.models.segmentation.fcn_resnet50(num_classes=2)
self.criterion = nn.CrossEntropyLoss()
self.evaluator = SMAPIoUMetric()
def forward(self, x):
return self.net(x)
def training_step(self, batch, batch_nb):
img, mask = batch
img = img.float()
mask = mask.long()
out = self.forward(img)["out"]
loss = self.criterion(out, mask)
self.log("loss", loss, prog_bar=True, sync_dist=True)
return loss
def validation_step(self, batch, batch_nb):
img, mask = batch
img = img.float()
mask = mask.long()
out = self.forward(img)["out"]
loss = self.criterion(out, mask)
probs = torch.softmax(out, dim=1)
preds = torch.argmax(probs, dim=1)
preds = preds.detach().cpu().numpy()
mask = mask.detach().cpu().numpy()
self.evaluator.process(input={"pred": preds, "gt": mask})
self.log(
"val_loss",
loss,
prog_bar=True,
sync_dist=True,
on_step=False,
on_epoch=True,
)
def on_validation_epoch_end(self) -> None:
metrics = self.evaluator.evaluate(0)
self.log(
f"val_high_vegetation_IoU",
metrics["high_vegetation__IoU"],
sync_dist=True,
)
self.log(f"val_mIoU", metrics["mIoU"], sync_dist=True)
def configure_optimizers(self):
return torch.optim.Adam(self.net.parameters(), lr=self.learning_rate)