def train_on_batch(model, batch, optimizer):
"""One train step on batch of MNIST data. Uses CrossEntropyLoss.
Parameters
----------
model : nn.Module
Model for MNIST classification.
batch : tuple
Tuple of images and labels
optimizer : torch Optimizer
Description of parameter `optimizer`.
Returns
-------
tuple: loss, accuracy
Both are numpy
"""
criterion = nn.CrossEntropyLoss()
images, labels = batch
outputs = model(images)
loss = criterion(outputs, labels)
# backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
return loss.cpu().detach().numpy(), st.accuracy(outputs.cpu(), labels.cpu())
def validate(model,val_loader,_log = logging.getLogger('validate')):
total = 0
total_accuracy = 0.0
total_loss = 0.0
model = model.eval()
_log.info('begin to validate with %d steps'%len(val_loader))
for images,labels in tqdm(val_loader):
with torch.no_grad():
images = images.cuda()
labels = labels.cuda()
criterion = nn.CrossEntropyLoss()
batch_size = images.shape[0]
predictions = model(images)
loss = criterion(predictions, labels)
total_loss += loss.cpu().detach().numpy() * batch_size
accuracy = st.accuracy(predictions.cpu(),labels.cpu())
total_accuracy += accuracy*batch_size
total += batch_size
total_accuracy /= total
total_loss /= total
model = model.train()
return total_loss, total_accuracy
def train_on_batch_resnet(model,batch,optimizer,isgpu = True):
if isinstance(model,nn.DataParallel):
resnet = model.module
else:
resnet = model
images,labels = batch
if isgpu:
images,labels = images.cuda(),labels.cuda()
# forward
predictions = resnet(images)
if isgpu:
criterion = nn.CrossEntropyLoss().cuda()
else:
criterion = nn.CrossEntropyLoss()
# print('predictions shape {}'.format(predictions.cpu().detach().numpy().shape))
# print('labels shape {}'.format(labels.shape))
loss = criterion(predictions,labels)
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss = loss.cpu().detach().numpy()
acc = st.accuracy(predictions.cuda(),labels.cuda())
return loss,acc
def train_on_batch(model,batch,optimizer,device = torch.device('cuda:0'),isgpu = True):
if isinstance(model,nn.DataParallel):
ode_model = model.module
else:
ode_model = model
images,labels = batch
if isgpu:
images = images.cuda()
labels = labels.cuda()
criterion = nn.CrossEntropyLoss()
# forward and inference
nfe_forward = ode_model.odefunc.nfe.item()
predictions = ode_model(images)
loss = criterion(predictions,labels)
# backward and optimizer
ode_model.odefunc.nfe.fill_(0)
optimizer.zero_grad()
loss.backward()
optimizer.step()
nfe_backward = ode_model.odefunc.nfe.item()
loss = loss.cpu().detach().numpy()
accuracy = st.accuracy(predictions.cpu(),labels.cpu())
return loss,accuracy,nfe_forward,nfe_backward
def classification_train_on_batch(model, batch, optimizer):
# batch is tuple containing (tensor of images, tensor of labels)
outputs = model(batch[0]) # forward pass
# compute loss
criterion = nn.CrossEntropyLoss()
loss = criterion(outputs, batch[1])
# backward pass and weight update
optimizer.zero_grad()
loss.backward()
optimizer.step()
# compute and return metrics
loss = loss.detach().cpu().numpy()
acc = st.accuracy(outputs, batch[1])
return loss, acc
def validate(model, val_loader, _log=logging.getLogger("validate")):
"""Find loss and accuracy on the given dataloader using the model.
Parameters
----------
model : nn.Module
Model for MNIST classification.
val_loader : DataLoader
Data over which to validate.
Returns
-------
tuple: val_loss, accuracy
Both are numpy
"""
with torch.no_grad():
model = model.eval()
val_loss = 0
accuracy = 0
total = 0
_log.info(f"Running validate with {len(val_loader)} steps")
for images, labels in tqdm(val_loader):
criterion = nn.CrossEntropyLoss()
batch_size = images.shape[0]
# forward pass
outputs = model(images)
loss = criterion(outputs, labels)
val_loss += loss*batch_size
accuracy += st.accuracy(outputs, labels)*batch_size
total += batch_size
val_loss /= total
accuracy /= total
model = model.train()
return val_loss.cpu().item(), accuracy
def train_on_batch(model, batch, optimizer):
"""One train step on batch of MNIST data. Uses CrossEntropyLoss.
Parameters
----------
model : nn.Module
Model for MNIST classification.
batch : tuple
Tuple of images and labels
optimizer : torch Optimizer
Description of parameter `optimizer`.
Returns
-------
tuple: loss, accuracy
Both are numpy
"""
if isinstance(model, nn.DataParallel):
ode_model = model.module
else:
ode_model = model
criterion = nn.CrossEntropyLoss()
images, labels = batch
ode_model.odefunc.nfe.fill_(0)
outputs = model(images)
nfe_forward = ode_model.odefunc.nfe.item()
loss = criterion(outputs, labels)
# backward and optimize
ode_model.odefunc.nfe.fill_(0)
optimizer.zero_grad()
loss.backward()
optimizer.step()
nfe_backward = ode_model.odefunc.nfe.item()
loss = loss.cpu().detach().numpy()
acc = st.accuracy(outputs.cpu(), labels.cpu())
return loss, acc, nfe_forward, nfe_backward
def classification_callback(model, val_loader, batch_metrics_dict):
with torch.no_grad(): # dont compute gradients
criterion = nn.CrossEntropyLoss()
model.eval() # eval mode
batches = len(val_loader)
loss = 0
acc = 0
for batch in val_loader:
outputs = model(batch[0])
loss += criterion(outputs, batch[1])
acc += st.accuracy(outputs, batch[1])
# find average loss and accuracy over whole vaildation set
loss /= batches
acc /= batches
model.train() # go back to train mode
# return metrics
return loss.cpu().numpy(), acc