# Show the training information
if batch % 500 == 0 or batch == len(val_loader):
acc = val_correct_cnt / val_total_cnt
ave_loss = val_total_loss / batch
print(
'Validation batch index: {}, val loss: {:.6f}, acc: {:.3f}'
.format(batch, ave_loss, acc))
validation_loss.append(ave_loss)
validation_acc.append(acc)
model.train()
# Save trained model
torch.save(model.state_dict(), './checkpoint/%s.pth' % model.name())
# Plot Learning Curve
# TODO
fig, axs = plt.subplots(nrows=2, ncols=2, constrained_layout=True)
axs[0, 0].plot(train_loss)
axs[0, 0].set_xlabel('epoch', fontsize=12)
axs[0, 0].set_ylabel('loss', fontsize=12)
axs[0, 0].set_title('Training Loss', fontsize=14)
axs[0, 1].plot(validation_loss)
axs[0, 1].set_xlabel('epoch', fontsize=12)
axs[0, 1].set_ylabel('loss', fontsize=12)
axs[0, 1].set_title('Validation Loss', fontsize=14)
axs[1, 0].plot(train_acc)
import torchvision.transforms as transforms
from model import ConvNet, Fully
from data import TestDataset
if __name__ == "__main__":
data_path, model_type, output = sys.argv[1], sys.argv[2], sys.argv[3]
if model_type == 'conv':
model = ConvNet()
elif model_type == 'fully':
model = Fully()
#######################################################################
# Modifiy this part to load your trained model
# TODO
model.load_state_dict(torch.load('./checkpoint/%s.pth' % model.name()))
#######################################################################
use_cuda = torch.cuda.is_available()
if use_cuda:
model.cuda()
model.eval()
# Load data
trans = transforms.Compose([
transforms.Grayscale(),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (1.0, ))
])
test_set = TestDataset(data_path, transform=trans)
print('Length of Testing Set:', len(test_set))
