dlmacedo / entropic-out-of-distribution-detection Goto Github PK

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    history = []

    for epoch in range(num_epochs):
        result = {}

        # each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0
            correct_females = 0
            correct_males = 0

            # iterate over data
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                for i, pred in enumerate(preds):
                    if labels[i] == 0 and labels[i] == pred: 
                      correct_females += 1
                    if labels[i] == 1 and labels[i] == pred:
                      correct_males += 1
            
            if phase == 'train':
                # record & update learning rate
                result['lrs'] = get_lr(optimizer)
                scheduler.step()

            # for printing to command line
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double().item() / dataset_sizes[phase]
            female_acc = correct_females / gender_dist[phase]['females']
            male_acc = correct_males / gender_dist[phase]['males']
            # for plotting
            result[phase+'_loss'] = epoch_loss
            result[phase+'_acc'] = epoch_acc
            result[phase+'_female_acc'] = female_acc
            result[phase+'_male_acc'] = male_acc

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                # Save if the model has best accuracy till now
                torch.save(model, data_dir+'_model_best.pt')

    torch.save(history, data_dir+'_history.pt')

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model, history


# Load pretrained ResNet Model and use seamless solution
import losses # from https://github.com/dlmacedo/entropic-out-of-distribution-detection/tree/master/losses
model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = losses.IsoMaxPlusLossFirstPart(num_ftrs, len(class_names))
model = model.to(device)
criterion = losses.IsoMaxPlusLossSecondPart()

# Observe that all parameters are being optimized
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) 

# # Freeze until 2nd block (included): https://raminnabati.com/post/002_adv_pytorch_freezing_layers/
ct = 0
for child in model.children():
  ct += 1
  if ct < 6:
      for param in child.parameters():
          param.requires_grad = False

# Decay LR by a factor of 0.5 every 5 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)


# Train
model, history = train_model(model, criterion, optimizer, exp_lr_scheduler, num_epochs=10)

logits = model(data)
probabilities = torch.nn.Softmax(dim=1)(logits)
if score_type == "MPS": # the maximum probability score
    soft_out = probabilities.max(dim=1)[0]
elif score_type == "ES": # the negative entropy score
    soft_out = (probabilities * torch.log(probabilities)).sum(dim=1)
elif score_type == "MDS": # the minimum distance score
    soft_out = logits.max(dim=1)[0]

  probabilities_at_targets=probabilities_for_training[range(distances.size(0)), targets]