def _train(self, epoch):
train_loss = MetricTracker()
for idx, batch in enumerate(self.train_loader):
image = batch['image'].to(self.device)
label = batch['label'].to(self.device)
self.optimizer.zero_grad()
out = self.model(image)
loss = self.criterion(out, label)
#print(f'loss: {loss}')
loss.backward()
self.optimizer.step()
self.lr_scheduler.step()
train_loss.update(loss.item(), out.shape[0])
print(f'#{idx} -- loss: {train_loss.avg:.4f}')
print(f'Training loss: {train_loss.avg:.4f}')
return {'loss': train_loss.avg}
def train(trainloader, model, optimizer, lossfunc, label_type, epoch,
use_cuda):
lossTracker = MetricTracker()
# set model to train mode
model.train()
# main training loop
for idx, data in enumerate(tqdm(trainloader, desc="training")):
numSample = data["image"].size(0)
# unpack sample
bands = data["image"]
if label_type == 'multi_label':
labels = data["label"]
else:
labels = (torch.max(data["label"], 1)[1]).type(torch.long)
# move data to gpu if model is on gpu
if use_cuda:
bands = bands.to(torch.device("cuda"))
labels = labels.to(torch.device("cuda"))
# reset gradients
optimizer.zero_grad()
# forward pass
logits = model(bands)
loss = lossfunc(logits, labels)
# backward pass
loss.backward()
optimizer.step()
#
lossTracker.update(loss.item(), numSample)
# train_writer.add_scalar("loss", lossTracker.avg, epoch)
wandb.log({'loss': lossTracker.avg, 'epoch': epoch})
print('Train loss: {:.6f}'.format(lossTracker.avg))
def _valid(self, epoch):
valid_loss = MetricTracker()
#valid_acc = MetricTracker()
self.model.eval()
# Iterate over data
for idx, sample in enumerate(self.valid_loader):
# get the inputs
image = sample['image'].to(self.device)
label = sample['label'].to(self.device)
# forward
out = self.model(image)
loss = self.criterion(out, label)
valid_loss.update(loss.item(), out.shape[0])
print(f'Valid Loss: {valid_loss.avg:.4f}')
self.model.train()
return {'loss': valid_loss.avg}