def accumulate(self, logits, targs):
with torch.no_grad():
semantic_logit_list = self.semantic_softmax_processor.split_logits_to_semantic_logits(
logits)
semantic_targets_tensor = self.semantic_softmax_processor.convert_targets_to_semantic_targets(
targs)
accuracy_list = []
accuracy_valid_list = []
result = 0
for i in range(
len(semantic_logit_list)): # scanning hirarchy_level_list
logits_i = semantic_logit_list[i]
targets_i = semantic_targets_tensor[:, i]
pred_i = logits_i.argmax(dim=-1)
ind_valid = (targets_i >= 0)
num_valids = torch.sum(ind_valid)
accuracy_valid_list.append(num_valids)
if num_valids > 0:
accuracy_list.append(
(pred_i[ind_valid] == targets_i[ind_valid]
).float().mean())
else:
accuracy_list.append(0)
result += accuracy_list[-1] * accuracy_valid_list[-1]
num_valids_total = sum(accuracy_valid_list)
result = result.detach()
num_valids_total = num_valids_total.detach()
if num_distrib() > 1:
result = reduce_tensor(result, num_distrib())
num_valids_total = reduce_tensor(num_valids_total, num_distrib())
self.total += result.item()
self.count += num_valids_total.item()
def validate_21k(val_loader, model):
print_at_master("starting validation")
model.eval()
top1 = AverageMeter()
top5 = AverageMeter()
with torch.no_grad():
for i, (input, target) in enumerate(val_loader):
# mixed precision
with autocast():
logits = model(input).float()
# measure accuracy and record loss
acc1, acc5 = accuracy(logits, target, topk=(1, 5))
if num_distrib() > 1:
acc1 = reduce_tensor(acc1, num_distrib())
acc5 = reduce_tensor(acc5, num_distrib())
torch.cuda.synchronize()
top1.update(acc1.item(), input.size(0))
top5.update(acc5.item(), input.size(0))
print_at_master("Validation results:")
print_at_master('Acc_Top1 [%] {:.2f}, Acc_Top5 [%] {:.2f} '.format(
top1.avg, top5.avg))
model.train()
def train_21k(model, train_loader, val_loader, optimizer, args):
# set loss
loss_fn = CrossEntropyLS(args.label_smooth)
# set scheduler
scheduler = lr_scheduler.OneCycleLR(optimizer,
max_lr=args.lr,
steps_per_epoch=len(train_loader),
epochs=args.epochs,
pct_start=0.1,
cycle_momentum=False,
div_factor=20)
# set scalaer
scaler = GradScaler()
# training loop
for epoch in range(args.epochs):
if num_distrib() > 1:
train_loader.sampler.set_epoch(epoch)
# train epoch
print_at_master("\nEpoch {}".format(epoch))
epoch_start_time = time.time()
for i, (input, target) in enumerate(train_loader):
with autocast(): # mixed precision
output = model(input)
loss = loss_fn(output, target) # note - loss also in fp16
model.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
scheduler.step()
epoch_time = time.time() - epoch_start_time
print_at_master(
"\nFinished Epoch, Training Rate: {:.1f} [img/sec]".format(
len(train_loader) * args.batch_size / epoch_time *
max(num_distrib(), 1)))
# validation epoch
validate_21k(val_loader, model)
def create_data_loaders(args):
data_path_train = os.path.join(args.data_path, 'imagenet21k_train')
train_transform = transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
CutoutPIL(cutout_factor=0.5),
RandAugment(),
transforms.ToTensor(),
])
data_path_val = os.path.join(args.data_path, 'imagenet21k_val')
val_transform = transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
])
train_dataset = ImageFolder(data_path_train, transform=train_transform)
val_dataset = ImageFolder(data_path_val, transform=val_transform)
print_at_master("length train dataset: {}".format(len(train_dataset)))
print_at_master("length val dataset: {}".format(len(val_dataset)))
sampler_train = None
sampler_val = None
if num_distrib() > 1:
sampler_train = torch.utils.data.distributed.DistributedSampler(
train_dataset)
sampler_val = OrderedDistributedSampler(val_dataset)
# Pytorch Data loader
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=sampler_train is None,
num_workers=args.num_workers,
pin_memory=True,
sampler=sampler_train)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
pin_memory=False,
sampler=sampler_val)
train_loader = PrefetchLoader(train_loader)
val_loader = PrefetchLoader(val_loader)
return train_loader, val_loader