Created
May 22, 2019 14:17
-
-
Save mirceast/65afe6134730e72794a922a8a4da3c86 to your computer and use it in GitHub Desktop.
Transfer Learning 3
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def train_model(model, criterion, optimizer, scheduler, num_epochs=25): | |
since = time.time() | |
best_model_wts = copy.deepcopy(model.state_dict()) | |
best_acc = 0.0 | |
epoch_time = [] # we'll keep track of the time needed for each epoch | |
for epoch in range(num_epochs): | |
epoch_start = time.time() | |
print('Epoch {}/{}'.format(epoch+1, num_epochs)) | |
print('-' * 10) | |
# Each epoch has a training and validation phase | |
for phase in ['train', 'val']: | |
if phase == 'train': | |
scheduler.step() | |
model.train() # Set model to training mode | |
else: | |
model.eval() # Set model to evaluate mode | |
running_loss = 0.0 | |
running_corrects = 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 training phase | |
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) | |
epoch_loss = running_loss / dataset_sizes[phase] | |
epoch_acc = running_corrects.double() / dataset_sizes[phase] | |
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_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()) | |
# Add the epoch time | |
t_epoch = time.time() - epoch_start | |
epoch_time.append(t_epoch) | |
print() | |
time_elapsed = time.time() - since | |
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) | |
print('Best val Acc: {:4f}'.format(best_acc)) | |
# Load best model weights | |
model.load_state_dict(best_model_wts) | |
return model, epoch_time |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment