if torch.cuda.is_available():
label = label.type(torch.cuda.LongTensor)
else:
label = label.type(torch.LongTensor)
for i, batch in enumerate(train_loader, 1):
global_count = global_count + 1
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
p = args.shot * args.way
data_shot, data_query = data[:p], data[p:]
logits = model(data_shot, data_query)
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)
print('epoch {}, train {}/{}, loss={:.4f} acc={:.4f}'.format(
epoch, i, len(train_loader), loss.item(), acc))
tl.add(loss.item())
ta.add(acc)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tl = tl.item()
ta = ta.item()
global_count = global_count + 1
if torch.cuda.is_available():
data, _ = [_.cuda() for _ in batch]
else:
data = batch[0]
p = args.shot * args.way
data_shot, data_query = data[:p], data[p:]
logits = model(data_shot, data_query) # KqN x KN x 1
# use logits to weights all labels, KN x N
prediction = torch.sum(
torch.mul(logits, label_support_onehot.unsqueeze(0)),
1) # KqN x N
# compute loss
loss = F.cross_entropy(prediction, label)
acc = count_acc(prediction, label)
writer.add_scalar('data/loss', float(loss), global_count)
writer.add_scalar('data/acc', float(acc), global_count)
print('epoch {}, train {}/{}, loss={:.4f} acc={:.4f}'.format(
epoch, i, len(train_loader), loss.item(), acc))
tl.add(loss.item())
ta.add(acc)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tl = tl.item()
ta = ta.item()
