def train(self):
"""The function for the meta-train phase."""
# Set the meta-train log
trlog = {}
trlog['args'] = vars(self.args)
trlog['train_loss'] = []
trlog['val_loss'] = []
trlog['train_acc'] = []
trlog['val_acc'] = []
trlog['max_acc'] = 0.0
trlog['max_acc_epoch'] = 0
# Set the timer
timer = Timer()
# Set global count to zero
global_count = 0
# Set tensorboardX
writer = SummaryWriter(comment=self.args.save_path)
# Generate the labels for train set of the episodes
label_shot = torch.arange(self.args.way).repeat(self.args.shot)
if torch.cuda.is_available():
label_shot = label_shot.type(torch.cuda.LongTensor)
else:
label_shot = label_shot.type(torch.LongTensor)
# Start meta-train
for epoch in range(1, self.args.max_epoch + 1):
# Update learning rate
self.lr_scheduler.step()
# Set the model to train mode
self.model.train()
# Set averager classes to record training losses and accuracies
train_loss_averager = Averager()
train_acc_averager = Averager()
# Generate the labels for test set of the episodes during meta-train updates
label = torch.arange(self.args.way).repeat(self.args.train_query)
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
# Using tqdm to read samples from train loader
tqdm_gen = tqdm.tqdm(self.train_loader)
for i, batch in enumerate(tqdm_gen, 1):
# Update global count number
global_count = global_count + 1
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
p = self.args.shot * self.args.way
data_shot, data_query = data[:p], data[p:]
# Output logits for model
logits = self.model((data_shot, label_shot, data_query))
# Calculate meta-train loss
loss = F.cross_entropy(logits, label)
# Calculate meta-train accuracy
acc = count_acc(logits, label)
# Write the tensorboardX records
writer.add_scalar('data/loss', float(loss), global_count)
writer.add_scalar('data/acc', float(acc), global_count)
# Print loss and accuracy for this step
tqdm_gen.set_description(
'Epoch {}, Loss={:.4f} Acc={:.4f}'.format(
epoch, loss.item(), acc))
# Add loss and accuracy for the averagers
train_loss_averager.add(loss.item())
train_acc_averager.add(acc)
# Loss backwards and optimizer updates
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the averagers
train_loss_averager = train_loss_averager.item()
train_acc_averager = train_acc_averager.item()
# Start validation for this epoch, set model to eval mode
self.model.eval()
# Set averager classes to record validation losses and accuracies
val_loss_averager = Averager()
val_acc_averager = Averager()
# Generate the labels for test set of the episodes during meta-val for this epoch
label = torch.arange(self.args.way).repeat(self.args.val_query)
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
# Print previous information
if epoch % 10 == 0:
print('Best Epoch {}, Best Val Acc={:.4f}'.format(
trlog['max_acc_epoch'], trlog['max_acc']))
# Run meta-validation
for i, batch in enumerate(self.val_loader, 1):
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
p = self.args.shot * self.args.way
data_shot, data_query = data[:p], data[p:]
logits = self.model((data_shot, label_shot, data_query))
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
val_loss_averager.add(loss.item())
val_acc_averager.add(acc)
# Update validation averagers
val_loss_averager = val_loss_averager.item()
val_acc_averager = val_acc_averager.item()
# Write the tensorboardX records
writer.add_scalar('data/val_loss', float(val_loss_averager), epoch)
writer.add_scalar('data/val_acc', float(val_acc_averager), epoch)
# Print loss and accuracy for this epoch
print('Epoch {}, Val, Loss={:.4f} Acc={:.4f}'.format(
epoch, val_loss_averager, val_acc_averager))
# Update best saved model
if val_acc_averager > trlog['max_acc']:
trlog['max_acc'] = val_acc_averager
trlog['max_acc_epoch'] = epoch
self.save_model('max_acc')
# Save model every 10 epochs
if epoch % 10 == 0:
self.save_model('epoch' + str(epoch))
# Update the logs
trlog['train_loss'].append(train_loss_averager)
trlog['train_acc'].append(train_acc_averager)
trlog['val_loss'].append(val_loss_averager)
trlog['val_acc'].append(val_acc_averager)
# Save log
torch.save(trlog, osp.join(self.args.save_path, 'trlog'))
if epoch % 10 == 0:
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(),
timer.measure(epoch / self.args.max_epoch)))
writer.close()
def train(self):
trlog = {}
trlog['args'] = vars(self.args)
trlog['train_loss'] = []
trlog['val_loss'] = []
trlog['train_acc'] = []
trlog['val_acc'] = []
trlog['max_acc'] = 0.0
trlog['max_acc_epoch'] = 0
timer = Timer()
global_count = 0
writer = SummaryWriter(comment=self.args.save_path)
for epoch in range(1, self.args.pre_max_epoch + 1):
self.lr_scheduler.step()
self.model.train()
self.model.mode = 'pre'
tl = Averager()
ta = Averager()
tqdm_gen = tqdm.tqdm(self.train_loader)
for i, batch in enumerate(tqdm_gen, 1):
global_count = global_count + 1
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
label = batch[1]
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
logits = self.model(data)
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
writer.add_scalar('data/loss', float(loss), global_count)
writer.add_scalar('data/acc', float(acc), global_count)
tqdm_gen.set_description(
'Epoch {}, Loss={:.4f} Acc={:.4f}'.format(
epoch, loss.item(), acc))
tl.add(loss.item())
ta.add(acc)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
tl = tl.item()
ta = ta.item()
self.model.eval()
self.model.mode = 'preval'
vl = Averager()
va = Averager()
label = torch.arange(self.args.way).repeat(self.args.val_query)
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
label_shot = torch.arange(self.args.way).repeat(self.args.shot)
if torch.cuda.is_available():
label_shot = label_shot.type(torch.cuda.LongTensor)
else:
label_shot = label_shot.type(torch.LongTensor)
print('Best Epoch {}, Best Val acc={:.4f}'.format(
trlog['max_acc_epoch'], trlog['max_acc']))
for i, batch in enumerate(self.val_loader, 1):
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
p = self.args.shot * self.args.way
data_shot, data_query = data[:p], data[p:]
logits = self.model((data_shot, label_shot, data_query))
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
vl.add(loss.item())
va.add(acc)
vl = vl.item()
va = va.item()
writer.add_scalar('data/val_loss', float(vl), epoch)
writer.add_scalar('data/val_acc', float(va), epoch)
print('Epoch {}, Val, Loss={:.4f} Acc={:.4f}'.format(
epoch, vl, va))
if va > trlog['max_acc']:
trlog['max_acc'] = va
trlog['max_acc_epoch'] = epoch
self.save_model('max_acc')
if epoch % 20 == 0:
self.save_model('epoch' + str(epoch))
trlog['train_loss'].append(tl)
trlog['train_acc'].append(ta)
trlog['val_loss'].append(vl)
trlog['val_acc'].append(va)
torch.save(trlog, osp.join(self.args.save_path, 'trlog'))
if epoch > self.args.pre_max_epoch - 2:
self.save_model('epoch-last')
torch.save(
self.optimizer.state_dict(),
osp.join(self.args.save_path, 'optimizer_latest.pth'))
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(), timer.measure(epoch / self.args.max_epoch)))
writer.close()
def train(self):
"""The function for the pre-train phase."""
# Set the pretrain log
trlog = {}
trlog['args'] = vars(self.args)
trlog['train_loss'] = []
trlog['val_loss'] = []
trlog['train_acc'] = []
trlog['val_acc'] = []
trlog['train_iou'] = []
trlog['val_iou'] = []
trlog['max_iou'] = 0.0
trlog['max_iou_epoch'] = 0
# Set the timer
timer = Timer()
# Set global count to zero
global_count = 0
# Set tensorboardX
writer = SummaryWriter(comment=self.args.save_path)
# Start pretrain
for epoch in range(1, self.args.pre_max_epoch + 1):
# Update learning rate
self.lr_scheduler.step()
# Set the model to train mode
self.model.train()
self.model.mode = 'train'
# Set averager classes to record training losses and accuracies
train_loss_averager = Averager()
train_acc_averager = Averager()
train_iou_averager = Averager()
# Using tqdm to read samples from train loader
tqdm_gen = tqdm.tqdm(self.train_loader)
for i, batch in enumerate(tqdm_gen, 1):
# Update global count number
global_count = global_count + 1
if torch.cuda.is_available():
data, label = [_.cuda() for _ in batch]
else:
data = batch[0]
label = batch[1]
# Output logits for model
logits = self.model(data)
# Calculate train loss
# CD loss is modified in the whole project to incorporate ony Cross Entropy loss. Modify as per requirement.
#loss = self.FL(logits, label) + self.CD(logits,label) + self.LS(logits,label)
loss = self.CD(logits, label)
# Calculate train accuracy
self._reset_metrics()
seg_metrics = eval_metrics(logits, label,
self.args.num_classes)
self._update_seg_metrics(*seg_metrics)
pixAcc, mIoU, _ = self._get_seg_metrics(
self.args.num_classes).values()
# Add loss and accuracy for the averagers
train_loss_averager.add(loss.item())
train_acc_averager.add(pixAcc)
train_iou_averager.add(mIoU)
# Print loss and accuracy till this step
tqdm_gen.set_description(
'Epoch {}, Loss={:.4f} Acc={:.4f} IOU={:.4f}'.format(
epoch, train_loss_averager.item(),
train_acc_averager.item() * 100.0,
train_iou_averager.item()))
# Loss backwards and optimizer updates
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the averagers
train_loss_averager = train_loss_averager.item()
train_acc_averager = train_acc_averager.item()
train_iou_averager = train_iou_averager.item()
writer.add_scalar('data/train_loss(Pre)',
float(train_loss_averager), epoch)
writer.add_scalar('data/train_acc(Pre)',
float(train_acc_averager) * 100.0, epoch)
writer.add_scalar('data/train_iou (Pre)',
float(train_iou_averager), epoch)
print(
'Epoch {}, Train: Loss={:.4f}, Acc={:.4f}, IoU={:.4f}'.format(
epoch, train_loss_averager, train_acc_averager * 100.0,
train_iou_averager))
# Start validation for this epoch, set model to eval mode
self.model.eval()
self.model.mode = 'val'
# Set averager classes to record validation losses and accuracies
val_loss_averager = Averager()
val_acc_averager = Averager()
val_iou_averager = Averager()
# Print previous information
if epoch % 1 == 0:
print('Best Val Epoch {}, Best Val IoU={:.4f}'.format(
trlog['max_iou_epoch'], trlog['max_iou']))
# Run validation
for i, batch in enumerate(self.val_loader, 1):
if torch.cuda.is_available():
data, labels, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
label = labels[0]
p = self.args.way * self.args.shot
data_shot, data_query = data[:p], data[p:]
label_shot, label = labels[:p], labels[p:]
par = data_shot, label_shot, data_query
logits = self.model(par)
# Calculate preval loss
#loss = self.FL(logits, label) + self.CD(logits,label) + self.LS(logits,label)
loss = self.CD(logits, label)
# Calculate val accuracy
self._reset_metrics()
seg_metrics = eval_metrics(logits, label, self.args.way)
self._update_seg_metrics(*seg_metrics)
pixAcc, mIoU, _ = self._get_seg_metrics(self.args.way).values()
val_loss_averager.add(loss.item())
val_acc_averager.add(pixAcc)
val_iou_averager.add(mIoU)
# Update validation averagers
val_loss_averager = val_loss_averager.item()
val_acc_averager = val_acc_averager.item()
val_iou_averager = val_iou_averager.item()
writer.add_scalar('data/val_loss(Pre)', float(val_loss_averager),
epoch)
writer.add_scalar('data/val_acc(Pre)',
float(val_acc_averager) * 100.0, epoch)
writer.add_scalar('data/val_iou (Pre)', float(val_iou_averager),
epoch)
# Print loss and accuracy for this epoch
print('Epoch {}, Val: Loss={:.4f} Acc={:.4f} IoU={:.4f}'.format(
epoch, val_loss_averager, val_acc_averager * 100.0,
val_iou_averager))
# Update best saved model
if val_iou_averager > trlog['max_iou']:
trlog['max_iou'] = val_iou_averager
trlog['max_iou_epoch'] = epoch
print("model saved in max_iou")
self.save_model('max_iou')
# Save model every 10 epochs
if epoch % 10 == 0:
self.save_model('epoch' + str(epoch))
# Update the logs
trlog['train_loss'].append(train_loss_averager)
trlog['train_acc'].append(train_acc_averager)
trlog['val_loss'].append(val_loss_averager)
trlog['val_acc'].append(val_acc_averager)
trlog['train_iou'].append(train_iou_averager)
trlog['val_iou'].append(val_iou_averager)
# Save log
torch.save(trlog, osp.join(self.args.save_path, 'trlog'))
if epoch % 1 == 0:
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(),
timer.measure(epoch / self.args.max_epoch)))
writer.close()
def train(self):
"""The function for the meta-train phase."""
# Set the meta-train log
#Change when resuming training
initial_epoch = 25
trlog = {}
trlog['args'] = vars(self.args)
trlog['train_loss'] = []
trlog['train_acc'] = []
trlog['train_iou'] = []
# Set the meta-val log
trlog['val_loss'] = []
trlog['val_acc'] = []
trlog['val_iou'] = []
trlog['max_iou'] = 0.2856
trlog['max_iou_epoch'] = 4
# Set the timer
timer = Timer()
# Set global count to zero
global_count = 0
# Set tensorboardX
writer = SummaryWriter(comment=self.args.save_path)
K = self.args.way + 1 #included Background as class
N = self.args.train_query
Q = self.args.test_query
# Start meta-train
for epoch in range(initial_epoch, self.args.max_epoch + 1):
print(
'----------------------------------------------------------------------------------------------------------------------------------------------------------'
)
# Update learning rate
self.lr_scheduler.step()
# Set the model to train mode
self.model.train()
# Set averager classes to record training losses and accuracies
train_loss_averager = Averager()
train_acc_averager = Averager()
train_iou_averager = Averager()
# Using tqdm to read samples from train loader
tqdm_gen = tqdm.tqdm(self.train_loader)
for i, batch in enumerate(tqdm_gen, 1):
# Update global count number
global_count = global_count + 1
if torch.cuda.is_available():
data, labels, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
labels = batch[1]
#print(data.shape)
#print(labels.shape)
p = K * N
im_train, im_test = data[:p], data[p:]
#Adjusting labels for each meta task
labels = downlabel(labels, K)
out_train, out_test = labels[:p], labels[p:]
'''
print(im_train.shape)
print(im_test.shape)
print(out_train.shape)
print(out_test.shape)
'''
if (torch.cuda.is_available()):
im_train = im_train.cuda()
im_test = im_test.cuda()
out_train = out_train.cuda()
out_test = out_test.cuda()
#Reshaping train set ouput
Ytr = out_train.reshape(-1)
Ytr = onehot(Ytr, K) #One hot encoding for loss
Yte = out_test.reshape(out_test.shape[0], -1)
if (torch.cuda.is_available()):
Ytr = Ytr.cuda()
Yte = Yte.cuda()
# Output logits for model
Gte = self.model(im_train, Ytr, im_test, Yte)
GteT = torch.transpose(Gte, 1, 2)
# Calculate meta-train loss
#loss = self.CD(GteT,Yte)
loss = self.FL(GteT, Yte)
#loss = self.LS(GteT,Yte)
self._reset_metrics()
# Calculate meta-train accuracy
seg_metrics = eval_metrics(GteT, Yte, K)
self._update_seg_metrics(*seg_metrics)
pixAcc, mIoU, _ = self._get_seg_metrics(K).values()
# Print loss and accuracy for this step
tqdm_gen.set_description(
'Epoch {}, Loss={:.4f} Acc={:.4f} IoU={:.4f}'.format(
epoch, loss.item(), pixAcc * 100.0, mIoU))
# Add loss and accuracy for the averagers
# Calculate the running averages
train_loss_averager.add(loss.item())
train_acc_averager.add(pixAcc)
train_iou_averager.add(mIoU)
# Loss backwards and optimizer updates
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the averagers
train_loss_averager = train_loss_averager.item()
train_acc_averager = train_acc_averager.item()
train_iou_averager = train_iou_averager.item()
#Adding to Tensorboard
writer.add_scalar('data/train_loss (Meta)',
float(train_loss_averager), epoch)
writer.add_scalar('data/train_acc (Meta)',
float(train_acc_averager) * 100.0, epoch)
writer.add_scalar('data/train_iou (Meta)',
float(train_iou_averager), epoch)
# Update best saved model if validation set is not present and save it
if (self.args.valdata == 'No'):
if train_iou_averager > trlog['max_iou']:
print("New Best!")
trlog['max_iou'] = train_iou_averager
trlog['max_iou_epoch'] = epoch
self.save_model('max_iou')
# Save model every 2 epochs
if epoch % 2 == 0:
self.save_model('epoch' + str(epoch))
# Update the logs
trlog['train_loss'].append(train_loss_averager)
trlog['train_acc'].append(train_acc_averager)
trlog['train_iou'].append(train_iou_averager)
if epoch % 1 == 0:
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(),
timer.measure(epoch / self.args.max_epoch)))
print('Epoch:{}, Average Loss: {:.4f}, Average mIoU: {:.4f}'.
format(epoch, train_loss_averager, train_iou_averager))
"""The function for the meta-val phase."""
if (self.args.valdata == 'Yes'):
# Start meta-val
# Set the model to val mode
self.model.eval()
# Set averager classes to record training losses and accuracies
val_loss_averager = Averager()
val_acc_averager = Averager()
val_iou_averager = Averager()
# Using tqdm to read samples from train loader
tqdm_gen = tqdm.tqdm(self.val_loader)
for i, batch in enumerate(tqdm_gen, 1):
# Update global count number
global_count = global_count + 1
if torch.cuda.is_available():
data, labels, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
labels = batch[1]
#print(data.shape)
#print(labels.shape)
p = K * N
im_train, im_test = data[:p], data[p:]
#Adjusting labels for each meta task
labels = downlabel(labels, K)
out_train, out_test = labels[:p], labels[p:]
'''
print(im_train.shape)
print(im_test.shape)
print(out_train.shape)
print(out_test.shape)
'''
if (torch.cuda.is_available()):
im_train = im_train.cuda()
im_test = im_test.cuda()
out_train = out_train.cuda()
out_test = out_test.cuda()
#Reshaping val set ouput
Ytr = out_train.reshape(-1)
Ytr = onehot(Ytr, K) #One hot encoding for loss
Yte = out_test.reshape(out_test.shape[0], -1)
if (torch.cuda.is_available()):
Ytr = Ytr.cuda()
Yte = Yte.cuda()
# Output logits for model
Gte = self.model(im_train, Ytr, im_test, Yte)
GteT = torch.transpose(Gte, 1, 2)
self._reset_metrics()
# Calculate meta-train accuracy
seg_metrics = eval_metrics(GteT, Yte, K)
self._update_seg_metrics(*seg_metrics)
pixAcc, mIoU, _ = self._get_seg_metrics(K).values()
# Print loss and accuracy for this step
tqdm_gen.set_description(
'Epoch {}, Val Loss={:.4f} Val Acc={:.4f} Val IoU={:.4f}'
.format(epoch, loss.item(), pixAcc * 100.0, mIoU))
# Add loss and accuracy for the averagers
# Calculate the running averages
val_loss_averager.add(loss.item())
val_acc_averager.add(pixAcc)
val_iou_averager.add(mIoU)
# Update the averagers
val_loss_averager = val_loss_averager.item()
val_acc_averager = val_acc_averager.item()
val_iou_averager = val_iou_averager.item()
#Adding to Tensorboard
writer.add_scalar('data/val_loss (Meta)',
float(val_loss_averager), epoch)
writer.add_scalar('data/val_acc (Meta)',
float(val_acc_averager) * 100.0, epoch)
writer.add_scalar('data/val_iou (Meta)',
float(val_iou_averager), epoch)
# Update best saved model
if val_iou_averager > trlog['max_iou']:
print("New Best (Validation)")
trlog['max_iou'] = val_iou_averager
trlog['max_iou_epoch'] = epoch
self.save_model('max_iou')
# Save model every 2 epochs
if epoch % 2 == 0:
self.save_model('epoch' + str(epoch))
# Update the logs
trlog['val_loss'].append(val_loss_averager)
trlog['val_acc'].append(val_acc_averager)
trlog['val_iou'].append(val_iou_averager)
if epoch % 1 == 0:
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(),
timer.measure(epoch / self.args.max_epoch)))
print(
'Epoch:{}, Average Val Loss: {:.4f}, Average Val mIoU: {:.4f}'
.format(epoch, val_loss_averager, val_iou_averager))
# Save log
torch.save(trlog, osp.join(self.args.save_path, 'trlog'))
print(
'----------------------------------------------------------------------------------------------------------------------------------------------------------'
)
writer.close()
def train(self):
"""The function for the pre-train phase."""
# Set the pretrain log
trlog = {}
trlog['args'] = vars(self.args)
trlog['train_loss'] = []
trlog['val_loss'] = []
trlog['train_acc'] = []
trlog['val_acc'] = []
trlog['max_acc'] = 0.0
trlog['max_acc_epoch'] = 0
# Set the timer
timer = Timer()
# Set global count to zero
global_count = 0
# Set tensorboardX
writer = SummaryWriter(comment=self.args.save_path)
# Start pretrain
for epoch in range(1, self.args.pre_max_epoch + 1):
# Set the model to train mode
print('Epoch {}'.format(epoch))
self.model.train()
self.model.mode = 'pre'
# Set averager classes to record training losses and accuracies
train_loss_averager = Averager()
train_acc_averager = Averager()
# Using tqdm to read samples from train loader
tqdm_gen = tqdm.tqdm(self.train_loader)
#for i, batch in enumerate(self.train_loader):
for i, batch in enumerate(tqdm_gen, 1):
# Update global count number
global_count = global_count + 1
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
label = batch[1]
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
logits = self.model(data)
loss = F.cross_entropy(logits, label)
# Calculate train accuracy
acc = count_acc(logits, label)
# Write the tensorboardX records
writer.add_scalar('data/loss', float(loss), global_count)
writer.add_scalar('data/acc', float(acc), global_count)
# Print loss and accuracy for this step
train_loss_averager.add(loss.item())
train_acc_averager.add(acc)
# Loss backwards and optimizer updates
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update the averagers
train_loss_averager = train_loss_averager.item()
train_acc_averager = train_acc_averager.item()
# start the original evaluation
self.model.eval()
self.model.mode = 'origval'
_, valid_results = self.val_orig(self.valset.X_val,
self.valset.y_val)
print('validation accuracy ', valid_results[0])
# Start validation for this epoch, set model to eval mode
self.model.eval()
self.model.mode = 'preval'
# Set averager classes to record validation losses and accuracies
val_loss_averager = Averager()
val_acc_averager = Averager()
# Generate the labels for test
label = torch.arange(self.args.way).repeat(self.args.val_query)
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
label_shot = torch.arange(self.args.way).repeat(self.args.shot)
if torch.cuda.is_available():
label_shot = label_shot.type(torch.cuda.LongTensor)
else:
label_shot = label_shot.type(torch.LongTensor)
# Run meta-validation
for i, batch in enumerate(self.val_loader, 1):
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
#data=data.float()
p = self.args.shot * self.args.way
data_shot, data_query = data[:p], data[p:]
logits = self.model((data_shot, label_shot, data_query))
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
val_loss_averager.add(loss.item())
val_acc_averager.add(acc)
# Update validation averagers
val_loss_averager = val_loss_averager.item()
val_acc_averager = val_acc_averager.item()
# Write the tensorboardX records
writer.add_scalar('data/val_loss', float(val_loss_averager), epoch)
writer.add_scalar('data/val_acc', float(val_acc_averager), epoch)
# Update best saved model
if val_acc_averager > trlog['max_acc']:
trlog['max_acc'] = val_acc_averager
trlog['max_acc_epoch'] = epoch
self.save_model('max_acc')
# Save model every 10 epochs
if epoch % 10 == 0:
self.save_model('epoch' + str(epoch))
# Update the logs
trlog['train_loss'].append(train_loss_averager)
trlog['train_acc'].append(train_acc_averager)
trlog['val_loss'].append(val_loss_averager)
trlog['val_acc'].append(val_acc_averager)
# Save log
torch.save(trlog, osp.join(self.args.save_path, 'trlog'))
if epoch % 10 == 0:
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(),
timer.measure(epoch / self.args.max_epoch)))
writer.close()
def train(self, trial):
"""The function for the meta-train phase."""
# Set the meta-train log
trlog = {}
trlog['args'] = vars(self.args)
trlog['train_loss'] = []
trlog['val_loss'] = []
trlog['train_acc'] = []
trlog['val_acc'] = []
trlog['max_acc'] = 0.0
trlog['max_acc_epoch'] = 0
# Set the timer
timer = Timer()
# Set global count to zero
global_count = 0
# Set tensorboardX
writer = SummaryWriter(comment=self.args.save_path)
# Generate the labels for train set of the episodes
label_shot = torch.arange(self.args.way).repeat(self.args.shot)
if torch.cuda.is_available():
label_shot = label_shot.type(torch.cuda.LongTensor)
else:
label_shot = label_shot.type(torch.LongTensor)
worstClasses = []
# Start meta-train
for epoch in range(1, self.args.max_epoch + 1):
# Update learning rate
self.lr_scheduler.step()
# Set the model to train mode
self.model.train()
# Set averager classes to record training losses and accuracies
train_loss_averager = Averager()
train_acc_averager = Averager()
# Generate the labels for test set of the episodes during meta-train updates
label = torch.arange(self.args.way).repeat(self.args.train_query)
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
# Using tqdm to read samples from train loader
tqdm_gen = tqdm.tqdm(self.train_loader)
for i, batch in enumerate(tqdm_gen, 1):
# Update global count number
global_count = global_count + 1
if torch.cuda.is_available():
data, targ = [_.cuda() for _ in batch]
else:
data, targ = batch
p = self.args.shot * self.args.way
data_shot, data_query = data[:p], data[p:]
# Output logits for model
if self.args.cross_att:
label_one_hot = self.one_hot(label).to(label.device)
ytest, cls_scores, logits = self.model(
(data_shot, label_shot, data_query),
ytest=label_one_hot)
pids = label_shot
loss = self.crossAttLoss(ytest, cls_scores, label, pids)
logits = logits[0]
else:
logits = self.model((data_shot, label_shot, data_query))
# Calculate meta-train loss
loss = F.cross_entropy(logits, label)
if self.args.distill_id:
teachLogits = self.teacher(
(data_shot, label_shot, data_query))
kl = F.kl_div(F.log_softmax(logits / self.args.kl_temp,
dim=1),
F.softmax(teachLogits / self.args.kl_temp,
dim=1),
reduction="batchmean")
loss = (kl * self.args.kl_interp * self.args.kl_temp *
self.args.kl_temp + loss *
(1 - self.args.kl_interp))
acc = count_acc(logits, label)
# Write the tensorboardX records
writer.add_scalar('data/loss', float(loss), global_count)
writer.add_scalar('data/acc', float(acc), global_count)
# Print loss and accuracy for this step
tqdm_gen.set_description(
'Epoch {}, Loss={:.4f} Acc={:.4f}'.format(
epoch, loss.item(), acc))
# Add loss and accuracy for the averagers
train_loss_averager.add(loss.item())
train_acc_averager.add(acc)
# Loss backwards and optimizer updates
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if self.args.hard_tasks:
if len(worstClasses) == self.args.way:
inds = self.train_sampler.hardBatch(worstClasses)
batch = [self.trainset[i][0] for i in inds]
data_shot, data_query = data[:p], data[p:]
logits = self.model(
(data_shot, label_shot, data_query))
loss = F.cross_entropy(logits, label)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
worstClasses = []
else:
error_mat = (logits.argmax(dim=1) == label).view(
self.args.train_query, self.args.way)
worst = error_mat.float().mean(dim=0).argmin()
worst_trueInd = targ[worst]
worstClasses.append(worst_trueInd)
# Update the averagers
train_loss_averager = train_loss_averager.item()
train_acc_averager = train_acc_averager.item()
# Start validation for this epoch, set model to eval mode
self.model.eval()
# Set averager classes to record validation losses and accuracies
val_loss_averager = Averager()
val_acc_averager = Averager()
# Generate the labels for test set of the episodes during meta-val for this epoch
label = torch.arange(self.args.way).repeat(self.args.val_query)
if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
# Print previous information
if epoch % 10 == 0:
print('Best Epoch {}, Best Val Acc={:.4f}'.format(
trlog['max_acc_epoch'], trlog['max_acc']))
# Run meta-validation
for i, batch in enumerate(self.val_loader, 1):
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
p = self.args.shot * self.args.way
data_shot, data_query = data[:p], data[p:]
if self.args.cross_att:
label_one_hot = self.one_hot(label).to(label.device)
ytest, cls_scores, logits = self.model(
(data_shot, label_shot, data_query),
ytest=label_one_hot)
pids = label_shot
loss = self.crossAttLoss(ytest, cls_scores, label, pids)
logits = logits[0]
else:
logits = self.model((data_shot, label_shot, data_query))
loss = F.cross_entropy(logits, label)
acc = count_acc(logits, label)
val_loss_averager.add(loss.item())
val_acc_averager.add(acc)
# Update validation averagers
val_loss_averager = val_loss_averager.item()
val_acc_averager = val_acc_averager.item()
# Write the tensorboardX records
writer.add_scalar('data/val_loss', float(val_loss_averager), epoch)
writer.add_scalar('data/val_acc', float(val_acc_averager), epoch)
# Print loss and accuracy for this epoch
print('Epoch {}, Val, Loss={:.4f} Acc={:.4f}'.format(
epoch, val_loss_averager, val_acc_averager))
# Update best saved model
if val_acc_averager > trlog['max_acc']:
trlog['max_acc'] = val_acc_averager
trlog['max_acc_epoch'] = epoch
self.save_model('max_acc')
# Update the logs
trlog['train_loss'].append(train_loss_averager)
trlog['train_acc'].append(train_acc_averager)
trlog['val_loss'].append(val_loss_averager)
trlog['val_acc'].append(val_acc_averager)
# Save log
torch.save(trlog, osp.join(self.args.save_path, 'trlog'))
if epoch % 10 == 0:
print('Running Time: {}, Estimated Time: {}'.format(
timer.measure(),
timer.measure(epoch / self.args.max_epoch)))
trial.report(val_acc_averager, epoch)
writer.close()