def main(model_config,
data_dir,
num_epochs=10,
batch_size=50,
learning_rate=0.001,
train_criterion=torch.nn.CrossEntropyLoss,
model_optimizer=torch.optim.Adam,
data_augmentations=None,
save_model_str=None):
"""
Training loop for configurableNet.
:param model_config: network config (dict)
:param data_dir: dataset path (str)
:param num_epochs: (int)
:param batch_size: (int)
:param learning_rate: model optimizer learning rate (float)
:param train_criterion: Which loss to use during training (torch.nn._Loss)
:param model_optimizer: Which model optimizer to use during trainnig (torch.optim.Optimizer)
:param data_augmentations: List of data augmentations to apply such as rescaling.
(list[transformations], transforms.Composition[list[transformations]], None)
If none only ToTensor is used
:return:
"""
# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if data_augmentations is None:
# You can add any preprocessing/data augmentation you want here
data_augmentations = transforms.ToTensor()
elif isinstance(type(data_augmentations), list):
data_augmentations = transforms.Compose(data_augmentations)
elif not isinstance(data_augmentations, transforms.Compose):
raise NotImplementedError
train_dataset = K49(data_dir, True, data_augmentations)
test_dataset = K49(data_dir, False, data_augmentations)
# Make data batch iterable
# Could modify the sampler to not uniformly random sample
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
model = torchModel(model_config,
input_shape=(train_dataset.channels,
train_dataset.img_rows,
train_dataset.img_cols),
num_classes=train_dataset.n_classes).to(device)
total_model_params = np.sum(p.numel() for p in model.parameters())
# instantiate optimizer
optimizer = model_optimizer(model.parameters(), lr=learning_rate)
# instantiate training criterion
train_criterion = train_criterion().to(device)
logging.info('Generated Network:')
summary(model, (train_dataset.channels, train_dataset.img_rows,
train_dataset.img_cols),
device='cuda' if torch.cuda.is_available() else 'cpu')
# Train the model
for epoch in range(num_epochs):
logging.info('#' * 50)
logging.info('Epoch [{}/{}]'.format(epoch + 1, num_epochs))
train_score, train_loss = model.train_fn(optimizer, train_criterion,
train_loader, device)
logging.info('Train accuracy %f', train_score)
test_score = model.eval_fn(test_loader, device)
logging.info('Test accuracy %f', test_score)
if save_model_str:
# Save the model checkpoint can be restored via "model = torch.load(save_model_str)"
if os.path.exists(save_model_str):
save_model_str += '_'.join(time.ctime())
torch.save(model.state_dict(), save_model_str)
def f_nn(params):
'''
lets say that you are making your own model from scratch,
you could do something like this but be sure of the shapes that you get in(:number of inchannels)
and also the the shape you output
if params['choice']['layers']== 'two':
self.fc1 = nn.Conv2d(channels, reduction, kernel_size=1, padding=0)
# calling the model function here with obove paramters
'''
model = torchModel()
model.to(device)
print('Params testing: ', params)
batch_size = int(params['batch_size'])
data_augmentations = transforms.ToTensor()
data_dir = '../data'
train_dataset = K49(data_dir, True, data_augmentations)
test_dataset = K49(data_dir, False, data_augmentations)
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
Train_dataset_loader = train_loader
Test_dataset_loader = test_loader
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),
lr=params['learning_rate'])
print('chossen learning rate', params['learning_rate'])
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
epochs = params['epochs']
steps = 0
train_losses, test_losses = [], []
for e in range(epochs):
correct = 0
average_precision = []
running_loss = 0
model.train()
exp_lr_scheduler.step()
for images, labels in Train_dataset_loader:
images, labels = Variable(images), Variable(labels)
images, labels = images.to(device), labels
optimizer.zero_grad()
log_ps = model(images)
loss = criterion(log_ps, labels.to(device))
loss.backward()
optimizer.step()
running_loss += loss.item(
) # calculate loss for batch wise and add it to the previous value
else:
test_loss = 0
accuracy = 0
total = 0
# Turn off gradients for validation, saves memory and computations
with torch.no_grad():
model.eval()
for images, labels in Test_dataset_loader:
images, labels = Variable(images), Variable(labels)
images, labels = images.to(device), labels.to(device)
ps = model(images)
test_loss += criterion(ps, labels.to(device))
_, predicted = torch.max(ps.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
train_losses.append(running_loss / len(Train_dataset_loader))
test_losses.append(test_loss / len(Test_dataset_loader))
if e == epochs - 1:
print("Epoch: {}/{}.. ".format(e + 1, epochs),
"Training Loss: {:.3f}.. ".format(train_losses[-1]),
"Test Loss: {:.3f}.. ".format(test_losses[-1]),
"Test Accuracy: {:.3f}".format(correct / total))
print('Accuracy of the network on the 10000 test images: %d %%' %
(100 * correct / total))
import matplotlib.pyplot as plt
plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)
loss = test_loss / len(Test_dataset_loader)
return loss.detach().item()
def compute(self, config, budget, working_directory, *args, **kwargs):
"""
Simple example for a compute function using a feed forward network.
It is trained on the MNIST dataset.
The input parameter "config" (dictionary) contains the sampled configurations passed by the bohb optimizer
"""
batch_size = int(config['batch_size'])
# Make data batch iterable
# Could modify the sampler to not uniformly random sample
self.train_loader = DataLoader(dataset=self.train_dataset,
batch_size=batch_size,
shuffle=True)
self.test_loader = DataLoader(dataset=self.test_dataset,
batch_size=batch_size,
shuffle=False)
# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = torchModel(config,
input_shape=(self.train_dataset.channels,
self.train_dataset.img_rows,
self.train_dataset.img_cols),
num_classes=self.train_dataset.n_classes).to(device)
total_model_params = np.sum(p.numel() for p in model.parameters())
# instantiate optimizer
if config['optimizer'] == 'Adam':
optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=config['lr'],
momentum=config['sgd_momentum'])
# instantiate training criterion
if config['train_criterion'] == 'cross_entropy':
train_criterion = torch.nn.CrossEntropyLoss()
else:
train_criterion = torch.nn.MSELoss()
logging.info('Generated Network:')
summary(model,
(self.train_dataset.channels, self.train_dataset.img_rows,
self.train_dataset.img_cols),
device='cuda' if torch.cuda.is_available() else 'cpu')
# Train the model
for epoch in range(int(budget)):
logging.info('#' * 50)
logging.info('Epoch [{}/{}]'.format(epoch + 1, int(budget)))
self.adjust_learning_rate(optimizer, epoch, config)
train_score, train_loss = model.train_fn(optimizer,
train_criterion,
self.train_loader, device)
logging.info('Train accuracy %f', train_score)
test_score = model.eval_fn(self.test_loader, device)
logging.info('Test accuracy %f', test_score)
return ({
'loss':
1 - (test_score / 100), # remember: HpBandSter always minimizes!
'info': {
'test accuracy': test_score,
'train accuracy': train_score,
'total_model_params': total_model_params
}
})