def create_datasets(self): transform = transforms.Compose([ transforms.RandomCrop(self.crop_size), self.net.transform, ]) loaders = [] for dataset in [self.source_dataset, self.target_dataset]: rot_dataset = RotationWrapper(dataset, transform=transform) loader = DataLoader(rot_dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=4) loaders.append(loader) self.train_dispenser = JointDispenser(*loaders) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: rot_dataset = RotationWrapper(dataset, transform=transform) loader = DataLoader(rot_dataset, batch_size=1, shuffle=True, num_workers=4) val_loaders.append(loader) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], }
def create_datasets(self): transform = transforms.Resize(1024) crop_transform = self.net.transform loaders = [] for dataset in [self.source_dataset, self.target_dataset]: grid_dataset = GridRegressionWrapper(dataset, stride=self.stride, transform=transform, crop_transform=crop_transform) loader = DataLoader(grid_dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=4) loaders.append(loader) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: grid_dataset = GridRegressionWrapper(dataset, stride=self.stride, transform=transform, crop_transform=crop_transform) loader = DataLoader(grid_dataset, batch_size=1, shuffle=True, num_workers=4) val_loaders.append(loader) self.train_dispenser = JointDispenser(*loaders) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], }
def create_datasets(self): transform = transforms.Compose([ self.net.transform, ]) target_transform = transforms.Compose([ to_tensor_raw, ]) collate_fn = AugmentCollate(crop=self.crop_size, flip=True) train_loaders = [] for train_dataset in [self.source_dataset, self.target_dataset]: dataset = TransformWrapper(train_dataset, transform=transform, target_transform=target_transform) train_loader = DataLoader(dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=8, collate_fn=collate_fn, drop_last=True) train_loaders.append(train_loader) self.train_dispenser = JointDispenser(*train_loaders) val_transform = transforms.Compose([ transforms.Resize((1024, 2048)), self.net.transform, ]) val_target_transform = transforms.Compose([ transforms.Resize((1024, 2048), interpolation=Image.NEAREST), to_tensor_raw, ]) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: dataset = TransformWrapper( dataset, transform=val_transform, target_transform=val_target_transform, ) loader = DataLoader( dataset, batch_size=torch.cuda.device_count(), num_workers=4, ) val_loaders.append(loader) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], }
def create_datasets(self): crop_transform = self.net.transform loaders = [] for dataset in [self.source_dataset, self.target_dataset]: grid_dataset = ContinuousGridRegressionWrapper( dataset, crop_size=self.crop_size, top=self.top, left=self.left, bottom=self.bottom, right=self.right, normalize=self.normalize, crop_transform=crop_transform ) loader = DataLoader( grid_dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=4, drop_last=True ) loaders.append(loader) self.train_dispenser = JointDispenser(*loaders) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: grid_dataset = ContinuousGridRegressionWrapper( dataset, crop_size=self.crop_size, top=self.top, left=self.left, bottom=self.bottom, right=self.right, normalize=self.normalize, crop_transform=crop_transform ) loader = DataLoader( grid_dataset, batch_size=4, shuffle=True, num_workers=4 ) val_loaders.append(loader) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], }
class Rotation: def __init__(self, net, source_dataset, target_dataset, source_val_dataset, target_val_dataset, *, batch_size, name='rotation', crop_size=400): self.net = net self.source_dataset = source_dataset self.target_dataset = target_dataset self.source_val_dataset = source_val_dataset self.target_val_dataset = target_val_dataset self.batch_size = batch_size self.name = name self.crop_size = crop_size self.loss_fn = nn.NLLLoss() self.create_datasets() self.create_head() def create_datasets(self): transform = transforms.Compose([ transforms.RandomCrop(self.crop_size), self.net.transform, ]) loaders = [] for dataset in [self.source_dataset, self.target_dataset]: rot_dataset = RotationWrapper(dataset, transform=transform) loader = DataLoader(rot_dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=4) loaders.append(loader) self.train_dispenser = JointDispenser(*loaders) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: rot_dataset = RotationWrapper(dataset, transform=transform) loader = DataLoader(rot_dataset, batch_size=1, shuffle=True, num_workers=4) val_loaders.append(loader) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], } def create_head(self): self.head = RotationHead(self.net.out_dim) self.net.attach_head(self.name, self.head) def _predict_batch(self, im): n, r, c, h, w = im.size() im = im.view(n * r, c, h, w).cuda() preds = self.net(im, task=self.name) preds = preds.view(n * r, 4) return preds def step(self): im, label = self.train_dispenser.next_batch() label = label.view(-1).cuda() preds = self._predict_batch(im) loss = self.loss_fn(preds, label) return loss def eval(self): self.net.eval() results = {} for domain, loader in self.val_loaders.items(): correct = 0 total = 0 for im, label in loader: with torch.no_grad(): label = label.view(-1).cuda() logits = self._predict_batch(im) preds = logits.max(dim=1)[1] correct += preds.eq(label).sum().item() total += label.numel() accuracy = correct / total logging.info(f' {self.name}.{domain}: {accuracy}') results[f'{self.name}.{domain}'] = accuracy self.net.train() return results
class AdversarialAdaptation: def __init__(self, net, source_dataset, target_dataset, source_val_dataset, target_val_dataset, *, batch_size, crop_size, name='adversarial_adapt'): self.name = name self.net = net self.source_dataset = source_dataset self.target_dataset = target_dataset self.source_val_dataset = source_val_dataset self.target_val_dataset = target_val_dataset self.adv_opt = torch.optim self.batch_size = batch_size self.crop_size = crop_size self.loss_fn = nn.CrossEntropyLoss() self.create_datasets() self.create_head() def create_datasets(self): transform = transforms.Compose([ self.net.transform, ]) target_transform = transforms.Compose([ to_tensor_raw, ]) collate_fn = AugmentCollate(crop=self.crop_size, flip=True) train_loaders = [] for train_dataset in [self.source_dataset, self.target_dataset]: dataset = TransformWrapper(train_dataset, transform=transform, target_transform=target_transform) train_loader = DataLoader(dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=8, collate_fn=collate_fn, drop_last=True) train_loaders.append(train_loader) self.train_dispenser = JointDispenser(*train_loaders) val_transform = transforms.Compose([ transforms.Resize((1024, 2048)), self.net.transform, ]) val_target_transform = transforms.Compose([ transforms.Resize((1024, 2048), interpolation=Image.NEAREST), to_tensor_raw, ]) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: dataset = TransformWrapper( dataset, transform=val_transform, target_transform=val_target_transform, ) loader = DataLoader( dataset, batch_size=torch.cuda.device_count(), num_workers=4, ) val_loaders.append(loader) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], } def create_head(self): self.adv_dis = AdversarialDiscriminator(self.net.out_dim).cuda() self.adv_opt = torch.optim.SGD(self.adv_dis.parameters(), lr=0.007, momentum=0.9) def _predict_batch(self, im): im = im.cuda() preds = self.net(im, task=self.name) return preds def step(self): self.step_discriminator() return self.step_generator() def step_discriminator(self): im, _ = self.train_dispenser.next_batch() im = im.cuda() label = self.generate_label(im) self.adv_opt.zero_grad() with torch.no_grad(): _, fts = self.net(im, task=None) preds = self.adv_dis(fts) loss = self.loss_fn(preds, label) loss.backward() self.adv_opt.step() def step_generator(self): im, _ = self.train_dispenser.next_batch() im = im.cuda() label = self.generate_label(im, invert=True) _, fts = self.net(im, task=None) preds = self.adv_dis(fts) loss = self.loss_fn(preds, label) return loss def generate_label(self, im, invert=False): zeros = torch.zeros(im.size(0) // 2, device=im.device, dtype=torch.long) ones = torch.ones(im.size(0) // 2, device=im.device, dtype=torch.long) if invert: return torch.cat([ones, zeros], dim=0) return torch.cat([zeros, ones], dim=0) def eval(self): return {}
class GridRegression: def __init__(self, net, source_dataset, target_dataset, source_val_dataset, target_val_dataset, *, batch_size, stride=256, name='gridregression'): self.net = net self.source_dataset = source_dataset self.target_dataset = target_dataset self.source_val_dataset = source_val_dataset self.target_val_dataset = target_val_dataset self.batch_size = batch_size self.stride = stride self.name = name self.loss_fn = nn.MSELoss() self.create_datasets() self.create_head() def create_datasets(self): transform = transforms.Resize(1024) crop_transform = self.net.transform loaders = [] for dataset in [self.source_dataset, self.target_dataset]: grid_dataset = GridRegressionWrapper(dataset, stride=self.stride, transform=transform, crop_transform=crop_transform) loader = DataLoader(grid_dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=4) loaders.append(loader) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: grid_dataset = GridRegressionWrapper(dataset, stride=self.stride, transform=transform, crop_transform=crop_transform) loader = DataLoader(grid_dataset, batch_size=1, shuffle=True, num_workers=4) val_loaders.append(loader) self.train_dispenser = JointDispenser(*loaders) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], } def create_head(self): self.head = GridRegressionHead(self.net.out_dim) self.net.attach_head(self.name, self.head) def _predict_batch(self, im): n, g, c, h, w = im.size() im = im.view(n * g, c, h, w).cuda() preds = self.net(im, task=self.name) preds = preds.view(n * g, 2) return preds def step(self): im, label = self.train_dispenser.next_batch() label = label.view(-1, 2).cuda() preds = self._predict_batch(im) loss = self.loss_fn(preds, label) return loss def eval(self): self.net.eval() results = {} for domain, loader in self.val_loaders.items(): correct = 0 total = 0 for im, label in loader: with torch.no_grad(): label = label.view(-1, 2).cuda() preds = self._predict_batch(im) preds = preds.round() correct += preds.eq(label).all(dim=1).sum().item() total += label.size(0) accuracy = correct / total logging.info(f' {self.name}.{domain}: {accuracy}') results[f'{self.name}.{domain}'] = accuracy self.net.train() return results
class ContinuousGridRegression: def __init__(self, net, source_dataset, target_dataset, source_val_dataset, target_val_dataset, *, batch_size, crop_size=256, top=0, left=0, bottom=1024, right=2048, normalize=True, name='gridregression-cont'): self.net = net self.source_dataset = source_dataset self.target_dataset = target_dataset self.source_val_dataset = source_val_dataset self.target_val_dataset = target_val_dataset self.batch_size = batch_size self.crop_size = crop_size self.top = top self.left = left self.bottom = bottom self.right = right self.name = name self.normalize = normalize self.loss_fn = nn.MSELoss() self.create_datasets() self.create_head() def create_datasets(self): crop_transform = self.net.transform loaders = [] for dataset in [self.source_dataset, self.target_dataset]: grid_dataset = ContinuousGridRegressionWrapper( dataset, crop_size=self.crop_size, top=self.top, left=self.left, bottom=self.bottom, right=self.right, normalize=self.normalize, crop_transform=crop_transform ) loader = DataLoader( grid_dataset, batch_size=self.batch_size // 2, shuffle=True, num_workers=4, drop_last=True ) loaders.append(loader) self.train_dispenser = JointDispenser(*loaders) val_loaders = [] for dataset in [self.source_val_dataset, self.target_val_dataset]: grid_dataset = ContinuousGridRegressionWrapper( dataset, crop_size=self.crop_size, top=self.top, left=self.left, bottom=self.bottom, right=self.right, normalize=self.normalize, crop_transform=crop_transform ) loader = DataLoader( grid_dataset, batch_size=4, shuffle=True, num_workers=4 ) val_loaders.append(loader) self.val_loaders = { 'source': val_loaders[0], 'target': val_loaders[1], } def create_head(self): self.head = GridRegressionHead(self.net.out_dim) self.net.attach_head(self.name, self.head) def _predict_batch(self, im): im = im.cuda() preds = self.net(im, task=self.name) preds = preds.view(-1, 2) return preds def step(self): im, label = self.train_dispenser.next_batch() label = label.cuda() preds = self._predict_batch(im) loss = self.loss_fn(preds, label) return loss def eval(self): self.net.eval() results = {} for domain, loader in self.val_loaders.items(): errs = [] for im, label in loader: with torch.no_grad(): label = label.cuda() preds = self._predict_batch(im) err = self.loss_fn(preds, label) errs.append(err.item()) avg_err = np.mean(errs) logging.info(f' {self.name}.{domain}: {avg_err}') results[f'{self.name}.{domain}'] = avg_err self.net.train() return results