TorchvisionBenchmark Class — pytorch Architecture

Source Code

benchmarks/distributed/ddp/benchmark.py lines 175–194

class TorchvisionBenchmark(Benchmark):
    def __init__(self, device, distributed_backend, bucket_size, model):
        super().__init__(
            device,
            distributed_backend,
            bucket_size,
        )
        self.model = model

    def __str__(self):
        return f"{self.model} with batch size {self.batch_size}"

    def create_model(self):
        return torchvision.models.__dict__[self.model]().to(self.device)

    def generate_inputs(self):
        return [torch.rand([self.batch_size, 3, 224, 224], device=self.device)]

    def generate_target(self):
        return torch.tensor([1] * self.batch_size, dtype=torch.long, device=self.device)