U-ශුද්ධපුහුණු

19import numpy as np
20import torch
21import torch.utils.data
22import torchvision.transforms.functional
23from torch import nn
24
25from labml import lab, tracker, experiment, monit
26from labml.configs import BaseConfigs
27from labml_helpers.device import DeviceConfigs
28from labml_nn.unet.carvana import CarvanaDataset
29from labml_nn.unet import UNet

32class Configs(BaseConfigs):

39    device: torch.device = DeviceConfigs()

42    model: UNet

45    image_channels: int = 3

47    mask_channels: int = 1

50    batch_size: int = 1

52    learning_rate: float = 2.5e-4

55    epochs: int = 4

58    dataset: CarvanaDataset

60    data_loader: torch.utils.data.DataLoader

63    loss_func = nn.BCELoss()

65    sigmoid = nn.Sigmoid()

68    optimizer: torch.optim.Adam

70    def init(self):

72        self.dataset = CarvanaDataset(lab.get_data_path() / 'carvana' / 'train',
73                                      lab.get_data_path() / 'carvana' / 'train_masks')

75        self.model = UNet(self.image_channels, self.mask_channels).to(self.device)

78        self.data_loader = torch.utils.data.DataLoader(self.dataset, self.batch_size,
79                                                       shuffle=True, pin_memory=True)

81        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)

84        tracker.set_image("sample", True)

86    @torch.no_grad()
87    def sample(self, idx=-1):

93        x, _ = self.dataset[np.random.randint(len(self.dataset))]

95        x = x.to(self.device)

98        mask = self.sigmoid(self.model(x[None, :]))

100        x = torchvision.transforms.functional.center_crop(x, [mask.shape[2], mask.shape[3]])

102        tracker.save('sample', x * mask)

104    def train(self):

112        for _, (image, mask) in monit.mix(('Train', self.data_loader), (self.sample, list(range(50)))):

114            tracker.add_global_step()

116            image, mask = image.to(self.device), mask.to(self.device)

119            self.optimizer.zero_grad()

121            logits = self.model(image)

124            mask = torchvision.transforms.functional.center_crop(mask, [logits.shape[2], logits.shape[3]])

126            loss = self.loss_func(self.sigmoid(logits), mask)

128            loss.backward()

130            self.optimizer.step()

132            tracker.save('loss', loss)

134    def run(self):

138        for _ in monit.loop(self.epochs):

140            self.train()

142            tracker.new_line()

144            experiment.save_checkpoint()

147def main():

149    experiment.create(name='unet')

152    configs = Configs()

155    experiment.configs(configs, {})

158    configs.init()

161    experiment.add_pytorch_models({'model': configs.model})

164    with experiment.start():
165        configs.run()

169if __name__ == '__main__':
170    main()