Python setup_cuda Examples

Example #1

File: model_dl_test.py Project: LongJohnCoder/archai

import logging

import torch
import torch.backends.cudnn as cudnn
import numpy as np

from torch_testbed import utils, cifar10_models
from torch_testbed.timing import MeasureTime, print_all_timings, print_timing, get_timing, clear_timings
from torch_testbed.dataloader_dali import cifar10_dataloaders

utils.setup_logging()
utils.setup_cuda(42)

batch_size = 512
half = True

datadir = utils.full_path('~/torchvision_data_dir')
train_dl, test_dl = cifar10_dataloaders(datadir,
                                        train_batch_size=batch_size,
                                        test_batch_size=1024,
                                        cutout=0)

model = cifar10_models.resnet18().cuda()
lr, momentum, weight_decay = 0.025, 0.9, 3.0e-4
optim = torch.optim.SGD(model.parameters(),
                        lr,
                        momentum=momentum,
                        weight_decay=weight_decay)
crit = torch.nn.CrossEntropyLoss().cuda()

if half:

Example #2

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import torch
import numpy as np

from torch_testbed import utils, cifar10_models
from torch_testbed.timing import MeasureTime, print_all_timings, print_timing, get_timing

utils.create_logger()
utils.setup_cuda(42, local_rank=0)

batch_size = 512
half = True
model = cifar10_models.resnet18().cuda()
lr, momentum, weight_decay = 0.025, 0.9, 3.0e-4
optim = torch.optim.SGD(model.parameters(),
                        lr, momentum=momentum, weight_decay=weight_decay)
crit = torch.nn.CrossEntropyLoss().cuda()

if half:
    model = model.half()
    crit = crit.half()

@MeasureTime
def iter_dl(ts):
    i, d = 0, 0
    for x, l in ts:
        y = model(x)
        loss = crit(y, l)
        optim.zero_grad()

Example #3

File: model_dl_test.py Project: wayne9qiu/archai

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import logging

import torch
import torch.backends.cudnn as cudnn
import numpy as np

from torch_testbed import utils, cifar10_models
from torch_testbed.timing import MeasureTime, print_all_timings, print_timing, get_timing, clear_timings
from torch_testbed.dataloader_dali import cifar10_dataloaders


utils.create_logger()
utils.setup_cuda(42, 0)

batch_size = 512
half = True

datadir = utils.full_path('~/dataroot')
train_dl, test_dl = cifar10_dataloaders(datadir,
    train_batch_size=batch_size, test_batch_size=1024,
    cutout=0)

model = cifar10_models.resnet18().cuda()
lr, momentum, weight_decay = 0.025, 0.9, 3.0e-4
optim = torch.optim.SGD(model.parameters(),
                        lr, momentum=momentum, weight_decay=weight_decay)
crit = torch.nn.CrossEntropyLoss().cuda()