def test_same_input(average): preds = _input_miss_class.preds target = _input_miss_class.target preds_flat = torch.cat(list(preds), dim=0) target_flat = torch.cat(list(target), dim=0) mc = Accuracy(num_classes=NUM_CLASSES, average=average) for i in range(NUM_BATCHES): mc.update(preds[i], target[i]) class_res = mc.compute() func_res = accuracy(preds_flat, target_flat, num_classes=NUM_CLASSES, average=average) sk_res = sk_accuracy(target_flat, preds_flat) assert torch.allclose(class_res, torch.tensor(sk_res).float()) assert torch.allclose(func_res, torch.tensor(sk_res).float())
class BasicModel(pl.LightningModule): def __init__(self, model: torch.nn.Module, freeze_layers: List = [], loss: str = 'ce', max_epochs: int = 100, optimizer: str = 'adam', learning_rate: float = 1e-3, weight_decay: float = 1e-5, lr_scheduler: str = 'cosine', lr_decay_steps: List = [60, 80], lr_decay_rate: float = 0.1, final_lr: float = 0., nesterov: bool = False, **kwargs): super().__init__() self.save_hyperparameters(ignore='model') if freeze_layers: if 'ExcludeFC' in freeze_layers: for name, param in model.named_parameters(): if not any(layer in name for layer in ['classifier', 'fc']): param.requires_grad = False if hasattr(model, 'fc'): model.fc.weight.data.normal_(mean=0.0, std=0.01) model.fc.bias.data.zero_() elif hasattr(model, 'classifier'): model.classifier.weight.data.normal_(mean=0.0, std=0.01) model.classifier.bias.data.zero_() print('Kaiming Initialization of fc/classifier\n') else: for name, param in model.named_parameters(): if any(layer in name for layer in freeze_layers): param.requires_grad = False self.learning_rate = learning_rate self.model = model self.configure_loss() # metrics | Overall accuracy num_classes = kwargs.get('num_classes', None) self.train_acc = Accuracy(num_classes=num_classes) self.val_acc = Accuracy(num_classes=num_classes, compute_on_step=False) self.val_scores = ClassifyScore(num_classes, kwargs.get('classes', None)) self.test_acc = Accuracy(num_classes=num_classes, compute_on_step=False) self.val_verbose = False self.final_val = False def on_train_epoch_start(self) -> None: if self.hparams.freeze_layers: if self.current_epoch == 0: print('Freeze layer:') for layer in self.hparams.freeze_layers: if hasattr(self.model, layer): getattr(self.model, layer).eval() if self.current_epoch == 0: print(f'\t{layer}') def shared_step(self, batch, batch_idx): x, y = batch preds = self.model(x) loss = self.loss_function(preds, y) return loss, preds, y def training_step(self, batch, batch_idx): loss, preds, y = self.shared_step(batch, batch_idx) return {'loss': loss, 'preds': preds, 'target': y} def training_step_end(self, outputs): preds = outputs['preds'].softmax(-1) self.train_acc(preds, outputs['target']) self.log('train/loss', outputs['loss'], prog_bar=True) self.log('train/acc', self.train_acc, on_step=True, on_epoch=True, prog_bar=True) def validation_step(self, batch, batch_idx): _, preds, y = self.shared_step(batch, batch_idx) return {'preds': preds, 'target': y} def validation_step_end(self, outputs): preds = outputs['preds'].argmax(dim=1) self.val_acc.update(preds, outputs['target']) # self.log('val/loss', outputs['loss']) self.log('val/acc', self.val_acc, on_step=False, on_epoch=True, prog_bar=True) self.val_scores.update(outputs['target'].cpu().numpy(), preds.cpu().numpy()) def validation_epoch_end(self, outputs): scores = self.val_scores.get_scores() if self.val_verbose: print('\n\nval OA={:.2f} | AA={:.2f} | Kappa={:.2f}'.format( scores['OA'], scores['AA'], scores['Kappa'])) self.log('val/OA', scores['OA']) self.log('val/AA', scores['AA']) self.log('val/Kappa', scores['Kappa']) if self.final_val: self.val_scores.print_score(scores) self.val_scores.print_hist(scores['hist']) self.epoch_scores = scores self.val_scores.reset() # self.val_acc.reset() def test_step(self, batch, batch_idx): _, preds, y = self.shared_step(batch, batch_idx) return {'preds': preds, 'target': y} def test_step_end(self, outputs): preds = outputs['preds'].argmax(dim=1) self.test_acc(preds, outputs['target']) def test_epoch_end(self, outputs): print(self.test_metrics.compute()) def configure_optimizers(self): params = self.parameters() params = list(filter(lambda p: p.requires_grad, params)) weight_decay = self.hparams.get('weight_decay', 0) if self.hparams.optimizer == 'adam': optimizer = torch.optim.Adam(params, lr=self.learning_rate, weight_decay=weight_decay) elif self.hparams.optimizer == 'sgd': optimizer = torch.optim.SGD(params, lr=self.learning_rate, momentum=0.9, weight_decay=self.hparams.weight_decay, nesterov=self.hparams.nesterov) if self.hparams.lr_scheduler is None: return optimizer else: if self.hparams.lr_scheduler == 'step': scheduler = lrs.MultiStepLR(optimizer, self.hparams.lr_decay_steps, gamma=self.hparams.lr_decay_rate) elif self.hparams.lr_scheduler == 'cosine': scheduler = lrs.CosineAnnealingLR( optimizer, T_max=self.hparams.max_epochs, eta_min=self.hparams.final_lr) else: raise ValueError('Invalid lr_scheduler type!') return [optimizer], [scheduler] def configure_loss(self): loss = self.hparams.loss.lower() if loss == 'ce': criterion = nn.CrossEntropyLoss elif loss == 'bce': criterion = nn.BCEWithLogitsLoss elif loss == 'mse': criterion = nn.MSELoss elif loss == 'l1': criterion = nn.L1Loss elif loss == 'mlsm': criterion = nn.MultiLabelSoftMarginLoss else: raise ValueError("Invalid Loss Type!") self.loss_function = criterion()