Jetson Xavier NX上でSimSiamとCIFAR-10で表現学習をやってみた(5) facebookreseach github train関数の変更

こちらは「Jetson Xavier NX上でSimSiamとCIFAR-10で表現学習をやってみた facebookresearch Transformクラスを新規作成(4)」の続きになります

technoxs-stacker.hatenablog.com

目次

スポンサーリンク

1.実行環境

Jetson Xavier NX
ubuntu18.04
docker
python3.x
pytorch

->Jetson Xavier NX上におけるpytrorch環境構築は以下でやってますので、ご参考までに(^^)/

technoxs-stacker.hatenablog.com

2.コード変更

def train(train_loader, model, criterion, optimizer, epoch, args, gpu):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4f')
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses],
        prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    model.train()

    end = time.time()
    # for i, (images, _) in enumerate(train_loader):
    for i, ((y1, y2), _) in enumerate(train_loader, start=epoch * len(train_loader)):
        # measure data loading time
        data_time.update(time.time() - end)
        y1 = y1.cuda(gpu, non_blocking=True)
        y2 = y2.cuda(gpu, non_blocking=True)

        # if args.gpu is not None:
        #     images[0] = images[0].cuda(args.gpu, non_blocking=True)
        #     images[1] = images[1].cuda(args.gpu, non_blocking=True)

        # compute output and loss
        # p1, p2, z1, z2 = model(x1=images[0], x2=images[1])
        p1, p2, z1, z2 = model(x1=y1, x2=y2)
        loss = -(criterion(p1, z2).mean() + criterion(p2, z1).mean()) * 0.5

        # losses.update(loss.item(), images[0].size(0))
        losses.update(loss.item(), y1.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            progress.display(i)

※(6)へ続きます

参考

github.com

スポンサーリンク