def objective(trial: optuna.Trial):
model_cfg = {
"input_size": [32, 32],
"input_channel": 3,
"depth_multiple": 1.0,
"width_multiple": 1.0,
}
conv_type = trial.suggest_categorical("conv_type", ["Conv", "DWConv"])
kernel_size = trial.suggest_int("kernel_size", 3, 7, step=2)
n_channel_01 = trial.suggest_int("n_channel_01", 8, 64, step=8)
n_channel_02 = trial.suggest_int("n_channel_02", 8, 128, step=8)
linear_activation = trial.suggest_categorical("linear_activation",
["ReLU", "SiLU"])
n_channel_03 = trial.suggest_int("n_channel_03", 64, 256, step=8)
n_channel_04 = trial.suggest_int("n_channel_04", 32, 128, step=8)
n_repeat = trial.suggest_int("n_repeat", 1, 3)
backbone = [
[-1, n_repeat, conv_type, [n_channel_01, kernel_size, 1]],
[-1, 1, "MaxPool", [2]],
[-1, n_repeat, conv_type, [int(n_channel_02), kernel_size, 1]],
[-1, 1, "MaxPool", [2]],
[-1, 1, "Flatten", []],
[-1, 1, "Linear", [n_channel_03, linear_activation]],
[-1, 1, "Linear", [n_channel_04, linear_activation]],
[-1, 1, "Linear", [10]],
]
model_cfg.update({"backbone": backbone})
model = Model(model_cfg, verbose=True).to(device)
batch_size = trial.suggest_int("batch_size", 8, 256)
epochs = trial.suggest_int("epochs", 5, 20)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
sampler=subset_sampler)
test_loader = DataLoader(test_dataset, batch_size=batch_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
trainer = TorchTrainer(model, criterion, optimizer, device=device)
trainer.train(train_loader,
n_epoch=epochs,
test_dataloader=test_loader)
test_loss, test_accuracy = trainer.test(test_loader)
return test_loss
def get_trainer(path: str) -> Tuple[Model, TorchTrainer]:
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
model = Model(path, verbose=True)
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
trainer = TorchTrainer(model, criterion, optimizer, device=device)
return model, trainer
def test_custom_module(verbose: bool = False):
model = Model(
os.path.join("tests", "test_configs", "custom_module_model.yaml"),
verbose=verbose,
)
assert model(torch.rand(1, 3, 32, 32)).shape == torch.Size([1, 10])
assert count_model_params(model) == 138568
def validate_export_model(
model: Model, trainer: TorchTrainer,
test_loader: DataLoader) -> Tuple[bool, float, float]:
model.eval()
test_loss, test_accuracy = trainer.test(test_loader)
# model.fuse().eval()
model.export().eval()
test_loss_fused, test_accuracy_fused = trainer.test(test_loader)
is_loss_close = (np.isclose(test_loss, test_loss_fused, rtol=0.001)
or test_loss_fused < test_loss)
is_accuracy_close = (np.isclose(
test_accuracy, test_accuracy_fused, rtol=0.001)
or test_accuracy_fused > test_accuracy)
return (is_loss_close and is_accuracy_close), test_accuracy, test_loss
def test_profiler():
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
model = Model(os.path.join("tests", "test_configs", "example.yaml"),
verbose=True)
model.to(device)
profiler = model.profile(n_run=100, batch_size=32, input_size=(224, 224))
mac = profiler.get_macs(verbose=True)
print(f"Total MACs: {mac:,.0f}")
profiler.print_result(sort_by_rank=True)
assert mac == 1616970
assert profiler.result_array.sum() != 0
def test_model_train():
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
model = Model(os.path.join("tests", "test_configs", "example.yaml"),
verbose=True)
batch_size = 16
epochs = 1
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = datasets.CIFAR10("./data/cifar10",
train=True,
download=True,
transform=preprocess)
test_dataset = datasets.CIFAR10("./data/cifar10",
train=False,
download=True,
transform=preprocess)
subset_sampler = SubsetRandomSampler(np.arange(0, len(train_dataset), 2))
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
sampler=subset_sampler)
test_loader = DataLoader(test_dataset, batch_size=batch_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
trainer = TorchTrainer(model, criterion, optimizer, device=device)
trainer.train(train_loader, n_epoch=epochs, test_dataloader=test_loader)
test_loss, test_accuracy = trainer.test(test_loader)
print(test_loss, test_accuracy)
assert test_accuracy > 0.45 and test_loss < 1.5
def test_tta_model():
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
epochs = 1
model = Model(os.path.join("tests", "test_configs", "example.yaml"),
verbose=True)
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
train_loader, test_loader = prepare_cifar10()
aug_funcs = [aug_flip_ud, aug_flip_lr, aug_flip_lrud, aug_random_scale]
for epoch in range(epochs):
seen = 0
tta_correct = 0
running_loss = 0
pbar = tqdm(enumerate(train_loader))
for batch, (data, labels) in pbar:
data, labels = data.to(device), labels.to(device)
aug_idx = np.random.randint(len(aug_funcs) + 1)
if aug_idx < len(aug_funcs):
data = aug_funcs[aug_idx](data)
optimizer.zero_grad()
tta_out = model(data)
# tta_out = model(data, augment_func=aug_funcs)
# tta_out = F.softmax(tta_out, dim=-1).sum(dim=0)
tta_loss = criterion(tta_out, labels)
tta_loss.backward()
optimizer.step()
_, tta_predicted = torch.max(tta_out, 1)
seen += labels.size(0)
tta_correct += (tta_predicted == labels).sum().item()
running_loss += tta_loss.item()
pbar.set_description(
f"[Train]: [{epoch + 1:03d}] "
f"Loss: {(running_loss / (batch + 1)):.7f}, "
f"Accuracy: {(tta_correct / seen) * 100:.2f}%")
pbar.close()
with torch.no_grad():
seen = 0
tta_correct = 0
model_correct = 0
tta_running_loss = 0
model_running_loss = 0
pbar = tqdm(enumerate(test_loader))
for batch, (data, labels) in pbar:
data, labels = data.to(device), labels.to(device)
tta_out = model(data, augment_func=aug_funcs)
# tta_out = model(data, augment_func=aug_random_scale)
# tta_out = tta_out.sum(dim=0)
tta_out = F.softmax(tta_out, dim=-1).sum(dim=0)
tta_out = F.softmax(tta_out, dim=-1)
model_out = model(data)
tta_loss = criterion(tta_out, labels)
model_loss = criterion(model_out, labels)
_, tta_predicted = torch.max(tta_out, 1)
_, model_predicted = torch.max(model_out, 1)
seen += labels.size(0)
tta_correct += (tta_predicted == labels).sum().item()
model_correct += (model_predicted == labels).sum().item()
tta_running_loss += tta_loss.item()
model_running_loss += model_loss.item()
pbar.set_description(
f"[Test]: [{epoch + 1:03d}] "
f"Loss(TTA, No-TTA): {(tta_running_loss / (batch + 1)):.7f}, {(model_running_loss / (batch + 1)):.7f}, "
f"Accuracy(TTA, No-TTA): {(tta_correct / seen) * 100:.2f}%, {(model_correct / seen) * 100:.2f}%"
)