|
|
|
|
|
|
|
|
|
import optuna |
|
import gc |
|
import torch |
|
from torch.utils import data |
|
import torchvision.datasets as datasets |
|
from torchvision import transforms |
|
from torch import nn |
|
import random |
|
import my_utils as mu |
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
|
|
train_trans = transforms.Compose([ |
|
|
|
|
|
transforms.RandomPerspective(distortion_scale=0.6, p=0.4), |
|
transforms.GaussianBlur(kernel_size=(5, 11), sigma=(0.1, 0.2)), |
|
transforms.RandomRotation(degrees=(-8, 8)), |
|
transforms.ToTensor(), |
|
|
|
]) |
|
|
|
test_trans = transforms.Compose([ |
|
|
|
transforms.ToTensor(), |
|
|
|
]) |
|
|
|
|
|
|
|
trans = [transforms.ToTensor()] |
|
trans = transforms.Compose(trans) |
|
|
|
cifar_trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_trans) |
|
cifar_testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=test_trans) |
|
|
|
|
|
|
|
|
|
|
|
|
|
print("Read dataset and create dataloaders - 5%") |
|
|
|
|
|
def SpatialAveragePool(X): |
|
return torch.mean(X, dim=[2,3]) |
|
|
|
|
|
def init_weights(m): |
|
if type(m) == nn.Linear or type(m) == nn.Conv2d: |
|
torch.nn.init.xavier_uniform_(m.weight) |
|
|
|
|
|
loss = nn.CrossEntropyLoss() |
|
|
|
|
|
class MakiNet(nn.Module): |
|
def __init__(self, conv_arch, num_classes, dropout_rate=0.0001): |
|
super(MakiNet, self).__init__() |
|
self.out_classes = num_classes |
|
self.conv_arch = conv_arch |
|
k= 0 |
|
for i, (num_conv, in_channels, out_channels) in enumerate(conv_arch): |
|
self.add_module(f"maki_block{i}", MakiBlock(num_conv, in_channels, out_channels, dropout_rate)) |
|
|
|
|
|
k = out_channels |
|
print(str(k) + " parameters") |
|
|
|
self.classifier = nn.Sequential( |
|
nn.Flatten(), |
|
nn.Linear(k, 75), |
|
nn.Dropout(p=dropout_rate), |
|
nn.ReLU(), |
|
nn.Linear(75, 25), |
|
nn.Dropout(p=dropout_rate), |
|
nn.ReLU(), |
|
nn.Linear(25, num_classes) |
|
) |
|
|
|
def forward(self, x): |
|
out = x |
|
|
|
for i in range(len(self.conv_arch)): |
|
out = self._modules[f"maki_block{i}"](out) |
|
out = SpatialAveragePool(out) |
|
out = self.classifier(out) |
|
return out |
|
|
|
|
|
class MakiBlock(nn.Module): |
|
def __init__(self, num_conv, input_channels, output_channels, dropout_rate): |
|
super(MakiBlock, self).__init__() |
|
self.num_convs = num_conv |
|
self.linear = nn.Linear(input_channels, num_conv, bias=False) |
|
self.relu = nn.ReLU() |
|
|
|
|
|
self.dropout = nn.Dropout(p=dropout_rate) |
|
|
|
for i in range(num_conv): |
|
|
|
self.add_module(f"conv{i}", nn.Conv2d(input_channels, output_channels, kernel_size=5, padding=1, stride=1)) |
|
|
|
self.add_module(f"batch_norm{i}", nn.BatchNorm2d(output_channels)) |
|
|
|
|
|
def forward(self, x): |
|
|
|
x = x.to(device) |
|
|
|
avg_out = SpatialAveragePool(x) |
|
avg_out = avg_out.to(device) |
|
|
|
lin_out = self.linear(avg_out) |
|
lin_out - self.dropout(lin_out) |
|
|
|
a = self.relu(lin_out) |
|
|
|
|
|
total_output = [] |
|
for j in range(self.num_convs): |
|
|
|
out = self._modules[f"conv{j}"](x) |
|
out = self._modules[f"batch_norm{j}"](out) |
|
out = self.dropout(out) |
|
out = self.relu(out) |
|
|
|
s = a[:, j].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) * out |
|
total_output.append(s) |
|
|
|
total_output = torch.stack(total_output, dim=0) |
|
out = torch.sum(total_output, dim=0) |
|
return torch.Tensor(out) |
|
|
|
print("Create the model - 40%") |
|
|
|
|
|
|
|
def train_model(trail, study, device=device): |
|
gc.collect() |
|
torch.cuda.empty_cache() |
|
gpu_memory = torch.cuda.memory_allocated(device='cuda:0') |
|
print(f"GPU memory allocated: {gpu_memory}") |
|
|
|
try: |
|
batch_size = trail.suggest_int("batch_size", 32, 256) |
|
train_iter = data.DataLoader(cifar_trainset, batch_size, shuffle=True) |
|
test_iter = data.DataLoader(cifar_testset, batch_size, shuffle=True) |
|
|
|
|
|
dropout_rate = 0.15 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
num_conv = 3 |
|
|
|
|
|
|
|
|
|
|
|
|
|
model_arch = [ |
|
[3, 3, 120], |
|
[3, 120, 100], |
|
[3, 100, 80], |
|
] |
|
net = MakiNet(model_arch, 10, dropout_rate=dropout_rate) |
|
net.to(device) |
|
|
|
|
|
net.apply(init_weights) |
|
|
|
lr = trail.suggest_float("lr", 1e-5, 9e-1, log=True) |
|
|
|
wd=0 |
|
optimizer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=wd) |
|
loss = nn.CrossEntropyLoss() |
|
timer = mu.Timer() |
|
|
|
num_epochs = trail.suggest_int("num_epochs", 20, 50) |
|
|
|
metric = mu.Accumulator(3) |
|
train_loss = 0 |
|
train_acc = 0 |
|
for epoch in range(num_epochs): |
|
print(f"Epoch: {epoch}/ {num_epochs}") |
|
for i, (X, y) in enumerate(train_iter): |
|
timer.start() |
|
net.train() |
|
optimizer.zero_grad() |
|
X, y = X.to(device), y.to(device) |
|
y_hat = net(X) |
|
l = loss(y_hat, y) |
|
l.backward() |
|
optimizer.step() |
|
with torch.no_grad(): |
|
metric.add(l*X.shape[0], mu.accuracy(y_hat, y), X.shape[0]) |
|
timer.stop() |
|
train_loss, train_acc = metric[0]/metric[2], metric[1]/metric[2] |
|
if (i+1) % 50 == 0: |
|
print(f'batch {i+1}, train loss {train_loss:.3f}, train acc {train_acc:.3f}') |
|
test_acc = mu.evaluate_accuracy_gpu(net, test_iter) |
|
print(f'Test Accuracy for epoch {epoch+1} is {test_acc:.3f}') |
|
|
|
if epoch == 5 and test_acc <= 0.1: |
|
|
|
raise optuna.exceptions.TrialPruned() |
|
|
|
test_acc = mu.evaluate_accuracy_gpu(net, test_iter) |
|
test_acc_delta = test_acc - train_acc |
|
if test_acc_delta < -0.25: |
|
|
|
raise optuna.exceptions.TrialPruned() |
|
|
|
try: |
|
if test_acc > study.best_trials[0].values[0]: |
|
torch.save(net.state_dict(), f"attempt6{test_acc:.3f}.params") |
|
except IndexError: |
|
print("No best trial yet.") |
|
torch.save(net.state_dict(), f"attempt6{test_acc:.3f}.params") |
|
except Exception as e: |
|
print("Exception occurred") |
|
return optuna.exceptions.OptunaError(f"Exception occurred during training. {e}") |
|
return test_acc |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
study = optuna.load_study(study_name="attempt6", storage="sqlite:///db.sqlite3") |
|
|
|
|
|
|
|
study.optimize(lambda trial: train_model(trial, study, device), n_trials=1000) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|