Verified Commit 81bf349d authored by Max Ehrlich's avatar Max Ehrlich
Browse files

Add throughput tests

parent d4208bd3
import models
import torch
import torch.optim as optim
import argparse
import data
import time
device = torch.device('cuda')
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', choices=data.spatial_dataset_map.keys(), help='Dataset to use')
parser.add_argument('--batch_size', type=int, help='Batch size')
parser.add_argument('--data', help='Root folder for the dataset')
args = parser.parse_args()
spatial_dataset = data.jpeg_dataset_map[args.dataset](args.batch_size, args.data)
dataset_info = data.dataset_info[args.dataset]
spatial_model = models.SpatialResNet(dataset_info['channels'], dataset_info['classes'])
jpeg_model = models.JpegResNet(spatial_model, n_freqs=6).to(device)
optimizer = optim.Adam(jpeg_model.parameters())
t0 = time.perf_counter()
models.train(jpeg_model, device, spatial_dataset[0], optimizer, 0)
torch.cuda.synchronize()
t1 = time.perf_counter()
training_time = t1 - t0
t0 = time.perf_counter()
models.test(jpeg_model, device, spatial_dataset[1])
torch.cuda.synchronize()
t1 = time.perf_counter()
testing_time = t1 - t0
with open('{}_jpeg_throughput.csv'.format(args.dataset), 'w') as f:
f.write('Training, Testing\n')
f.write('{}, {}\n'.format(training_time / len(spatial_dataset[0]), testing_time / len(spatial_dataset[1])))
import models
import torch
import torch.optim as optim
import argparse
import data
import time
device = torch.device('cuda')
torch.backends.cudnn.enabled = False
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', choices=data.spatial_dataset_map.keys(), help='Dataset to use')
parser.add_argument('--batch_size', type=int, help='Batch size')
parser.add_argument('--data', help='Root folder for the dataset')
args = parser.parse_args()
spatial_dataset = data.spatial_dataset_map[args.dataset](args.batch_size, args.data)
dataset_info = data.dataset_info[args.dataset]
spatial_model = models.SpatialResNet(dataset_info['channels'], dataset_info['classes']).to(device)
optimizer = optim.Adam(spatial_model.parameters())
t0 = time.perf_counter()
models.train(spatial_model, device, spatial_dataset[0], optimizer, 0)
torch.cuda.synchronize()
t1 = time.perf_counter()
training_time = t1 - t0
t0 = time.perf_counter()
models.test(spatial_model, device, spatial_dataset[1])
torch.cuda.synchronize()
t1 = time.perf_counter()
testing_time = t1 - t0
with open('{}_spatial_throughput.csv'.format(args.dataset), 'w') as f:
f.write('Training, Testing\n')
f.write('{}, {}\n'.format(training_time / len(spatial_dataset[0]), testing_time / len(spatial_dataset[1])))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment