PyTorch フレームワークを使用した NNCF による量子化対応トレーニング#

def train(train_loader, model, criterion, optimizer, epoch): 
    batch_time = AverageMeter("Time", ":3.3f") 
    losses = AverageMeter("Loss", ":2.3f") 
    top1 = AverageMeter("Acc@1", ":2.2f") 
    top5 = AverageMeter("Acc@5", ":2.2f") 
    progress = ProgressMeter( 
        len(train_loader), 
        [batch_time, losses, top1, top5], 
        prefix="Epoch:[{}]".format(epoch), 
    ) 

    # トレーニング・モードに切り替え 
    model.train() 

    end = time.time() 
    for i, (images, target) in enumerate(train_loader): 
        images = images.to(device) 
        target = target.to(device) 

        # 出力を計算 
        output = model(images) 
        loss = criterion(output, target) 

        # 精度を測定し損失を記録 
        acc1, acc5 = accuracy(output, target, topk=(1, 5)) 
        losses.update(loss.item(), images.size(0)) 
        top1.update(acc1[0], images.size(0)) 
        top5.update(acc5[0], images.size(0)) 

        # 勾配を計算し opt ステップを実行 
        optimizer.zero_grad() 
        loss.backward() 
        optimizer.step() 

        # 経過時間を測定 
        batch_time.update(time.time() - end) 
        end = time.time() 

        print_frequency = 50 
        if i % print_frequency == 0: 
            progress.display(i)

def validate(val_loader, model, criterion): 
    batch_time = AverageMeter("Time", ":3.3f") 
    losses = AverageMeter("Loss", ":2.3f") 
    top1 = AverageMeter("Acc@1", ":2.2f") 
    top5 = AverageMeter("Acc@5", ":2.2f") 
    progress = ProgressMeter(len(val_loader), [batch_time, losses, top1, top5], prefix="Test: ") 

    # 評価モードに切り替え 
    model.eval() 

    with torch.no_grad(): 
        end = time.time() 
        for i, (images, target) in enumerate(val_loader): 
            images = images.to(device) 
            target = target.to(device) 

            # 出力を計算 
            output = model(images) 
            loss = criterion(output, target) 

            # 精度を測定し損失を記録 
            acc1, acc5 = accuracy(output, target, topk=(1, 5)) 
            losses.update(loss.item(), images.size(0)) 
            top1.update(acc1[0], images.size(0)) 
            top5.update(acc5[0], images.size(0)) 

            # 経過時間を測定 
            batch_time.update(time.time() - end) 
            end = time.time() 

            print_frequency = 10 
            if i % print_frequency == 0: 
                progress.display(i) 

        print(" * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}".format(top1=top1, top5=top5)) 
    return top1.avg

class AverageMeter(object): 
    """Computes and stores the average and current value""" 

    def __init__(self, name, fmt=":f"): 
        self.name = name 
        self.fmt = fmt 
        self.reset() 

    def reset(self): 
        self.val = 0 
        self.avg = 0 
        self.sum = 0 
        self.count = 0 

    def update(self, val, n=1): 
        self.val = val 
        self.sum += val * n 
        self.count += n 
        self.avg = self.sum / self.count 

    def __str__(self): 
        fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})" 
        return fmtstr.format(**self.__dict__) 

class ProgressMeter(object): 
    def __init__(self, num_batches, meters, prefix=""): 
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches) 
        self.meters = meters 
        self.prefix = prefix 

    def display(self, batch): 
        entries = [self.prefix + self.batch_fmtstr.format(batch)] 
        entries += [str(meter) for meter in self.meters] 
        print("\t".join(entries)) 

    def _get_batch_fmtstr(self, num_batches): 
        num_digits = len(str(num_batches // 1)) 
        fmt = "{:"+ str(num_digits) + "d}" 
        return "[" + fmt + "/" + fmt.format(num_batches) + "]" 

def accuracy(output, target, topk=(1,)): 
    """Computes the accuracy over the k top predictions for the specified values of k""" 
    with torch.no_grad(): 
        maxk = max(topk) 
        batch_size = target.size(0) 

        _, pred = output.topk(maxk, 1, True, True) 
        pred = pred.t() 
        correct = pred.eq(target.view(1, -1).expand_as(pred)) 

        res = [] 
        for k in topk: 
            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True) 
            res.append(correct_k.mul_(100.0 / batch_size)) 
        return res

num_classes = 200 # 200 は Tiny ImageNet 用、デフォルトは ImageNet 用の 1000 です 
init_lr = 1e-4 
batch_size = 128 
epochs = 4 

model = models.resnet18(pretrained=not pretrained_on_tiny_imagenet) 
# Tiny ImageNet のクラス数の最後の FC レイヤーを更新 
model.fc = nn.Linear(in_features=512, out_features=num_classes, bias=True) 
model.to(device) 

# データ読み込みコード 
train_dir = DATASET_DIR / "train" 
val_dir = DATASET_DIR / "val" 
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 

train_dataset = datasets.ImageFolder( 
    train_dir, 
    transforms.Compose( 
        [ 
            transforms.Resize(image_size), 
            transforms.RandomHorizontalFlip(), 
            transforms.ToTensor(), 
            normalize, 
        ] 
    ), 
) 
val_dataset = datasets.ImageFolder( 
    val_dir, 
    transforms.Compose( 
        [ 
            transforms.Resize(image_size), 
            transforms.ToTensor(), 
            normalize, 
        ] 
    ), 
) 

train_loader = torch.utils.data.DataLoader( 
    train_dataset, 
    batch_size=batch_size, 
    shuffle=True, 
    num_workers=0, 
    pin_memory=True, 
    sampler=None, 
) 

val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True) 

# 損失関数 (基準) とオプティマイザーを定義 
criterion = nn.CrossEntropyLoss().to(device) 
optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)

if pretrained_on_tiny_imagenet:
    # 
    # ** 警告: `torch.load` 機能は Python の pickling モジュールを使用します。 
    # このモジュールは、unpickle 中に任意のコードを実行するために使用できます。 
    # 信頼できるデータのみをロードします。 
    # 
    checkpoint = torch.load(str(fp32_pth_path), map_location="cpu") 
    model.load_state_dict(checkpoint["state_dict"], strict=True) 
    acc1_fp32 = checkpoint["acc1"] 
else: 
    best_acc1 = 0 
    # トレーニング・ループ 
    for epoch in range(0, epochs):
        # 単一のトレーニング・エポックを実行 
        train(train_loader, model, criterion, optimizer, epoch) 

        # 検証セットで評価 
        acc1 = validate(val_loader, model, criterion) 

        is_best = acc1 > best_acc1 
        best_acc1 = max(acc1, best_acc1) 

        if is_best: 
            checkpoint = {"state_dict": model.state_dict(), "acc1": acc1} 
            torch.save(checkpoint, fp32_pth_path) 
    acc1_fp32 = best_acc1 

print(f"Accuracy of FP32 model: {acc1_fp32:.3f}")

dummy_input = torch.randn(1, 3, image_size, image_size).to(device) 

ov_model = ov.convert_model(model, example_input=dummy_input, input=[1, 3, image_size, image_size]) 
ov.save_model(ov_model, fp32_ir_path, compress_to_fp16=False) 
print(f"FP32 model was exported to {fp32_ir_path}.")