class Test:
def __init__(self):
self.net_dict = NNState(mode='eval')
# Data Augmentation operations
img_transforms = transforms.Compose([
transforms.RandomRotation((-30, 30)),
transforms.RandomResizedCrop((64, 64), scale=(0.7, 1.0)),
transforms.ColorJitter(brightness=0.4,
contrast=0.3,
saturation=0.3,
hue=0.3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
self.eval_data = datasets.ImageFolder('./dataset_segmented/test',
transform=img_transforms)
def eval(self):
print('Evaluating...')
self.net_dict.net = self.net_dict.net.eval()
eval_loader = DataLoader(dataset=self.eval_data,
batch_size=self.net_dict.batch_size,
shuffle=False,
num_workers=0,
drop_last=False)
n_batch = len(eval_loader)
with torch.no_grad():
eval_loss_stack = self.net_dict.to_device(torch.Tensor())
correct = 0
total = 0
for i, batch in enumerate(eval_loader):
# forward propagation
inputs, labels = batch[0], batch[1]
inputs = self.net_dict.to_device(inputs)
labels = self.net_dict.to_device(labels)
# Forward
labels_hat = self.net_dict.net.forward(inputs)
_, predicted = torch.max(labels_hat.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
loss_batch = self.net_dict.criterion(labels_hat, labels)
eval_loss_stack = torch.cat(
(eval_loss_stack, loss_batch.unsqueeze(0)), 0)
print('Batch [%d/%d], Eval Loss: %.4f' %
(i + 1, n_batch, loss_batch))
eval_loss = torch.mean(eval_loss_stack)
print('*********************************')
print('=> Mean Evaluation Loss: %.3f' % eval_loss)
print('=> Accuracy of the network: %d %%' %
(100 * correct / total))
print('*********************************')
return eval_loss
class Train:
def __init__(self):
self.net_dict = NNState('train')
# Data Augmentation operations
img_transforms = transforms.Compose([
transforms.RandomRotation((-30, 30)),
transforms.RandomResizedCrop((64, 64), scale=(0.7, 1.0)),
transforms.ColorJitter(brightness=0.4,
contrast=0.3,
saturation=0.3,
hue=0.3),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
self.train_data = datasets.ImageFolder('./dataset_segmented/train',
transform=img_transforms)
print(self.train_data.class_to_idx)
self.eval_data = datasets.ImageFolder('./dataset_segmented/eval',
transform=img_transforms)
def train(self):
train_loader = DataLoader(dataset=self.train_data,
batch_size=self.net_dict.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
n_batch = len(train_loader)
for epoch_idx in range(self.net_dict.last_epoch + 1,
self.net_dict.n_epochs):
train_loss_buff = torch.Tensor()
train_loss_buff = self.net_dict.to_device(train_loss_buff)
print('\nEpoch [%d/%d]:' % (epoch_idx, self.net_dict.n_epochs))
t_start = time.time()
# update the network
for i, batch in enumerate(train_loader):
self.net_dict.optimiser.zero_grad()
inputs, labels = batch[0], batch[1]
inputs = self.net_dict.to_device(inputs)
labels = self.net_dict.to_device(labels)
# Forward
labels_hat = self.net_dict.net.forward(inputs)
loss = self.net_dict.criterion(labels_hat, labels)
# Backward
loss.backward()
# Optimise
self.net_dict.optimiser.step()
train_loss_buff = torch.cat(
(train_loss_buff, loss.reshape(1, 1)), 0)
if (i + 1) % 10 == 0:
print('[%d/%d], Itr [%d/%d], Loss: %.4f' %
(epoch_idx, self.net_dict.n_epochs, i, n_batch,
loss.item()))
# current_lr = self.optimiser.param_groups[0]['lr']
self.net_dict.lr_scheduler.step()
avg_train_loss = torch.mean(train_loss_buff)
print('=> Average training loss: %.4f' % avg_train_loss)
print('Training Duration: %.3fs' % (time.time() - t_start))
if (epoch_idx + 1) % 1 == 0:
eval_loss_mean = self.eval()
# Save model, and best model if qualified
delta_acc = self.net_dict.best_acc - eval_loss_mean
if delta_acc > 0:
self.net_dict.best_acc = eval_loss_mean
self.net_dict.save_ckpt(epoch_idx, delta_acc)
def eval(self):
print('Evaluating...')
self.net_dict.net = self.net_dict.net.eval()
eval_loader = DataLoader(dataset=self.eval_data,
batch_size=self.net_dict.batch_size,
shuffle=False,
num_workers=0,
drop_last=False)
n_batch = len(eval_loader)
with torch.no_grad():
eval_loss_stack = self.net_dict.to_device(torch.Tensor())
correct = 0
total = 0
for i, batch in enumerate(eval_loader):
# forward propagation
inputs, labels = batch[0], batch[1]
inputs = self.net_dict.to_device(inputs)
labels = self.net_dict.to_device(labels)
# Forward
labels_hat = self.net_dict.net.forward(inputs)
_, predicted = torch.max(labels_hat.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
loss_batch = self.net_dict.criterion(labels_hat, labels)
eval_loss_stack = torch.cat(
(eval_loss_stack, loss_batch.unsqueeze(0)), 0)
print('Batch [%d/%d], Eval Loss: %.4f' %
(i + 1, n_batch, loss_batch))
eval_loss = torch.mean(eval_loss_stack)
print('*********************************')
print('=> Mean Evaluation Loss: %.3f' % eval_loss)
print('=> Accuracy of the network: %d %%' %
(100 * correct / total))
print('*********************************')
return eval_loss
