Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More

How to use get_default_device method in Airtest

Best Python code snippet using Airtest

train.py

Source:train.py

...27def split_indices(n, val_pct):28    indices = np.random.permutation(n)29    bound = int(n * val_pct)30    return indices[bound:], indices[:bound]31def get_default_device():32    if cuda.is_available():33        return torch.device('cuda')34    else:35        return torch.device('cpu')  # incase cuda is unavaliable.36# move data and model to device. This operation may take while37def to_device(data, device):38    if isinstance(data, (list, tuple)):39        return [to_device(x, device) for x in data]  # breakdown element and40        # recursively put elements into device if the data is list or tuple.41    return data.to(device, non_blocking=True)  # This may return a pointer pointing to the data in GPU42# A wrapped dataloader that put data into gpu in BATCH-WISE fashion.43# This can save GPU dram.44class DeviceDataloader():45    # Constructor46    # @param data_loader A DataLoader instance that need to be operated.47    # @param device The device going to put on.48    def __init__(self, data_loader, device):49        self.data_loader = data_loader50        self.device = device51    # iterator function.52    def __iter__(self):53        for batch_tuple in self.data_loader:54            yield to_device(batch_tuple, self.device)  # do to_device. Yield is non-stop return,55    # length of self.data_loader56    def __len__(self):57        return len(self.data_loader)58class WR_Conv_Model(nn.Module):59    def accuracy(self, prediction, target):60        _, max_pos = torch.max(prediction, dim=1)61        return torch.tensor(torch.sum(max_pos == target).item() / len(max_pos))62    def __init__(self):63        super(WR_Conv_Model, self).__init__()64        self.conv_layer1 = nn.Conv2d(3, 6, kernel_size=3, padding=1, stride=1)65        # now BSx6x128x12866        self.conv_layer2 = nn.Conv2d(6, 12, kernel_size=3, padding=1, stride=1)67        # now BSx12x64x6468        self.conv_layer3 = nn.Conv2d(12, 24, kernel_size=3, padding=1, stride=1)69        # now BSx24x32x3270        self.conv_layer4 = nn.Conv2d(24, 48, kernel_size=3, padding=1, stride=1)71        # now BSx48x16x1672        self.conv_layer5 = nn.Conv2d(48, 96, kernel_size=3, padding=1, stride=1)73        # now BSx96x8x874        self.conv_layer6 = nn.Conv2d(96, 192, kernel_size=3, padding=1, stride=1)75        # now BSx192x4x476        # here data is flatten from dim=177        self.linear_layer1 = nn.Linear(192 * 4 * 4, 256)78        # now BSx25679        self.linear_layer2 = nn.Linear(256, 4)80        # now prediction is getted: BSx481    def forward(self, input_batch):82        out = self.conv_layer1(input_batch)83        out = nn.functional.relu(out)84        out = nn.functional.max_pool2d(out, kernel_size=2)85        out = self.conv_layer2(out)86        out = nn.functional.relu(out)87        out = nn.functional.max_pool2d(out, kernel_size=2)88        out = self.conv_layer3(out)89        out = nn.functional.relu(out)90        out = nn.functional.max_pool2d(out, kernel_size=2)91        out = self.conv_layer4(out)92        out = nn.functional.relu(out)93        out = nn.functional.max_pool2d(out, kernel_size=2)94        out = self.conv_layer5(out)95        out = nn.functional.relu(out)96        out = nn.functional.max_pool2d(out, kernel_size=2)97        out = self.conv_layer6(out)98        out = nn.functional.relu(out)99        out = nn.functional.max_pool2d(out, kernel_size=2)100        out = out.view(input_batch.size(0), -1)101        out = self.linear_layer1(out)102        out = nn.functional.relu(out)103        out = self.linear_layer2(out)104        return out105    def training_step(self, data_batch, target_batch):106        prediction = self(data_batch)107        loss = nn.functional.cross_entropy(prediction, target_batch)108        return loss109    def validation_step(self, data_batch, target_batch):110        prediction = self(data_batch)111        accuracy = self.accuracy(prediction, target_batch)112        loss = nn.functional.cross_entropy(prediction, target_batch)113        return {"validation loss": loss, "validation accuracy": accuracy}114# ---------------------------Prepare and load the data----------------------#115# open dataset116dataset = ImageFolder(ROOT_DIR + "/train", transform=ToTensor())117print("Dataset length: " + str(len(dataset)))118# split the indices into train and validation119train_ids, val_ids = split_indices(len(dataset), 0.1)  # 10% validation set120# transform data into dataloader121train_loader = DataLoader(dataset, BATCH_SIZE, sampler=SubsetRandomSampler(train_ids))122val_loader = DataLoader(dataset, BATCH_SIZE, sampler=SubsetRandomSampler(val_ids))123train_loader = DeviceDataloader(train_loader, get_default_device())124val_loader = DeviceDataloader(val_loader, get_default_device())125print("--------------Data transformed to: " + str(get_default_device()) + "--------------------")126# ----------------------------------END Prepare---------------------------------------#127# --------------------------------Construct model and optimizer-----------------------#128model = WR_Conv_Model()129to_device(model, get_default_device())130# optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)131optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)132# ---------------------------------END Construct---------------------------------------#133# --------------------------------Train------------------------------------------#134print("Here")135'''stage_1 = False136stage_2 = False137stage_3 = False'''138for i in range(EPOCH):139    # time0 = time.time()140    for data_batch, target_batch in train_loader:141        # print("Sample prediction: " + str(nn.functional.softmax(prediction, dim=1)[0]))142        loss = model.training_step(data_batch, target_batch)143        loss.backward()...

jit_models_test.py

Source:jit_models_test.py

...108        (torch.rand(3, 1, 4301) - 0.5) * 2,109    ),110)111def test_enhancement_model(small_model_params, model_def, test_data):112    device = get_default_device()113    params = small_model_params[model_def.__name__]114    model = model_def(**params)115    model = model.eval().to(device)116    # Random input uniformly distributed in [-1, 1]117    inputs = ((torch.rand(1, 2500, device=device) - 0.5) * 2,)118    traced = torch.jit.trace(model, inputs)119    assert_consistency(model=model, traced=traced, tensor=test_data.to(device))120@pytest.mark.parametrize("test_shape", [(2,), (3, 1)])121@pytest.mark.parametrize("matching_samples", [4701, 8800, 17001])122def test_dcunet_model(test_shape: Tuple, matching_samples):123    n_samples = 5010124    device = get_default_device()125    model = DCUNet(architecture="mini", fix_length_mode="pad").eval().to(device)126    # Random input uniformly distributed in [-1, 1]127    inputs = torch.rand(1, n_samples, device=device)128    traced = torch.jit.trace(model, (inputs,))129    test_data = torch.rand(*test_shape, matching_samples, device=device)130    assert_consistency(model=model, traced=traced, tensor=test_data.to(device))131@pytest.mark.parametrize(132    "model_def",133    (134        ConvTasNet,135        DPRNNTasNet,136        DPTNet,137        LSTMTasNet,138        SuDORMRFNet,139        SuDORMRFImprovedNet,140    ),141)142@pytest.mark.parametrize(143    "test_data",144    (145        (torch.rand(240) - 0.5) * 2,146        (torch.rand(1, 220) - 0.5) * 2,147        (torch.rand(3, 250) - 0.5) * 2,148        (torch.rand(1, 1, 301) - 0.5) * 2,149        (torch.rand(2, 1, 501) - 0.5) * 2,150    ),151)152def test_trace_bss_model(small_model_params, model_def, test_data):153    device = get_default_device()154    params = small_model_params[model_def.__name__]155    model = model_def(**params)156    model = model.eval().to(device)157    # Random input uniformly distributed in [-1, 1]158    inputs = ((torch.rand(1, 201, device=device) - 0.5) * 2,)159    traced = torch.jit.trace(model, inputs)160    assert_consistency(model=model, traced=traced, tensor=test_data.to(device))161def get_default_device():162    if torch.cuda.is_available():163        return "cuda"...

interface.py

Source:interface.py

2from pathlib import Path3import yaml4import torch5import os6def get_default_device():7    if torch.cuda.is_available():8        return "cuda"9    else:10        return "cpu"11def load_model(mel2wav_path, device=get_default_device()):12    """13    Args:14        mel2wav_path (str or Path): path to the root folder of dumped text2mel15        device (str or torch.device): device to load the model16    """17    root = Path(mel2wav_path)18    with open(root / "args.yml", "r") as f:19        args = yaml.load(f, Loader=yaml.FullLoader)20    netG = Generator(args.n_mel_channels, args.ngf, args.n_residual_layers).to(device)21    netG.load_state_dict(torch.load(root / "best_netG.pt", map_location=device))22    return netG23class MelVocoder:24    def __init__(25        self,26        path,27        device=get_default_device(),28        github=False,29        model_name="multi_speaker",30    ):31        self.fft = Audio2Mel().to(device)32        if github:33            netG = Generator(80, 32, 3).to(device)34            root = Path(os.path.dirname(__file__)).parent35            netG.load_state_dict(36                torch.load(root / f"models/{model_name}.pt", map_location=device)37            )38            self.mel2wav = netG39        else:40            self.mel2wav = load_model(path, device)41        self.device = device...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.