torchvision.datasets.ImageFolder(args.val, transform = data_transforms),
batch_size=args.batch_size, shuffle=True,
num_workers=args.worker
)
return train_loader, val_loader
train_loader, val_loader = get_dataloader()
classes = 10
net = LeNet(classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model to gpu
net.to(device)
print("Start Training...")
os.makedirs('./expr', exist_ok=True)
for epoch in range(1000):
net.train()
running_loss = 0.0
for i, data in enumerate(train_loader):
inputs, labels = data
inputs, labels = inputs.to(device, dtype=torch.float), labels.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
def main():
parser = argparse.ArgumentParser()
mode_group = parser.add_mutually_exclusive_group(required=True)
mode_group.add_argument("--train",
action="store_true",
help="To train the network.")
mode_group.add_argument("--test",
action="store_true",
help="To test the network.")
parser.add_argument("--epochs",
default=10,
type=int,
help="Desired number of epochs.")
parser.add_argument("--dropout",
action="store_true",
help="Whether to use dropout or not.")
parser.add_argument("--uncertainty",
action="store_true",
help="Use uncertainty or not.")
parser.add_argument("--dataset",
action="store_true",
help="The dataset to use.")
parser.add_argument("--outsample",
action="store_true",
help="Use out of sample test image")
uncertainty_type_group = parser.add_mutually_exclusive_group()
uncertainty_type_group.add_argument(
"--mse",
action="store_true",
help=
"Set this argument when using uncertainty. Sets loss function to Expected Mean Square Error."
)
uncertainty_type_group.add_argument(
"--digamma",
action="store_true",
help=
"Set this argument when using uncertainty. Sets loss function to Expected Cross Entropy."
)
uncertainty_type_group.add_argument(
"--log",
action="store_true",
help=
"Set this argument when using uncertainty. Sets loss function to Negative Log of the Expected Likelihood."
)
dataset_type_group = parser.add_mutually_exclusive_group()
dataset_type_group.add_argument(
"--mnist",
action="store_true",
help="Set this argument when using MNIST dataset")
dataset_type_group.add_argument(
"--emnist",
action="store_true",
help="Set this argument when using EMNIST dataset")
dataset_type_group.add_argument(
"--CIFAR",
action="store_true",
help="Set this argument when using CIFAR dataset")
dataset_type_group.add_argument(
"--fmnist",
action="store_true",
help="Set this argument when using FMNIST dataset")
args = parser.parse_args()
if args.dataset:
if args.mnist:
from mnist import dataloaders, label_list
elif args.CIFAR:
from CIFAR import dataloaders, label_list
elif args.fmnist:
from fashionMNIST import dataloaders, label_list
if args.train:
num_epochs = args.epochs
use_uncertainty = args.uncertainty
num_classes = 10
model = LeNet(dropout=args.dropout)
if use_uncertainty:
if args.digamma:
criterion = edl_digamma_loss
elif args.log:
criterion = edl_log_loss
elif args.mse:
criterion = edl_mse_loss
else:
parser.error(
"--uncertainty requires --mse, --log or --digamma.")
else:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.005)
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=7,
gamma=0.1)
device = get_device()
model = model.to(device)
model, metrics = train_model(model,
dataloaders,
num_classes,
criterion,
optimizer,
scheduler=exp_lr_scheduler,
num_epochs=num_epochs,
device=device,
uncertainty=use_uncertainty)
state = {
"epoch": num_epochs,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
}
if use_uncertainty:
if args.digamma:
torch.save(state, "./results/model_uncertainty_digamma.pt")
print("Saved: ./results/model_uncertainty_digamma.pt")
if args.log:
torch.save(state, "./results/model_uncertainty_log.pt")
print("Saved: ./results/model_uncertainty_log.pt")
if args.mse:
torch.save(state, "./results/model_uncertainty_mse.pt")
print("Saved: ./results/model_uncertainty_mse.pt")
else:
torch.save(state, "./results/model.pt")
print("Saved: ./results/model.pt")
elif args.test:
use_uncertainty = args.uncertainty
device = get_device()
model = LeNet()
model = model.to(device)
optimizer = optim.Adam(model.parameters())
if use_uncertainty:
if args.digamma:
checkpoint = torch.load(
"./results/model_uncertainty_digamma.pt")
if args.log:
checkpoint = torch.load("./results/model_uncertainty_log.pt")
if args.mse:
checkpoint = torch.load("./results/model_uncertainty_mse.pt")
else:
checkpoint = torch.load("./results/model.pt")
filename = "./results/rotate.jpg"
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.eval()
if args.outsample:
img = Image.open("./data/arka.jpg").convert('L').resize((28, 28))
img = TF.to_tensor(img)
img.unsqueeze_(0)
else:
a = iter(dataloaders['test'])
img, label = next(a)
rotating_image_classification(model,
img,
filename,
label_list,
uncertainty=use_uncertainty)
img = transforms.ToPILImage()(img[0][0])
test_single_image(model, img, label_list, uncertainty=use_uncertainty)
input_size = 32
input_size_h = 32
input_size_w = 40
number_classes = 2
in_channel = 1
int_version = 4
name = 'LeNet{}x{}_{}'.format(input_size_h, input_size_w, int_version)
#hardware setting
device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = 'cpu' # now it works only under 'cpu' settings
net = LeNet(in_channel, number_classes)
#if(device == 'cuda'):
net = net.to(device)
# load a pre-trained pyTorch model
#checkpoint = torch.load("./lenet32x40_ckpt_0_60.667504.pth")
checkpoint = torch.load("./ckpt_32x40_lenet_3_277_99.206645.pth")
net.load_state_dict(checkpoint['weight'])
# if u want to use cpu, then you need to do something
# net = net.to('cpu')
# input_ = input_.to('cpu')
net.eval()
input_ = torch.ones([1, in_channel, input_size_h, input_size_w])
input = input_.to(device)
# input=torch.ones([1,3,224,224])
def training(model_name, trainloader, validloader, input_channel=3, epochs=1, resume=True, self_define=True, only_print=False):
# load self defined or official net
assert model_name in ["LeNet", "VGG16", "ResNet", "DenseNet"]
if self_define:
if model_name == "LeNet":
net = LeNet(input_channel)
elif model_name == "VGG16":
net = VGG16(input_channel)
elif model_name == "ResNet":
net = ResNet(input_channel)
elif model_name == "DenseNet":
net = DenseNet(input_channel)
else:
if model_name == "LeNet":
net = LeNet(input_channel) # on official LeNet
elif model_name == "VGG16":
net = models.vgg16_bn(pretrained=False, num_classes=10)
elif model_name == "ResNet":
net = models.resnet50(pretrained=False, num_classes=10)
elif model_name == "DenseNet":
net = models.DenseNet(num_classes=10)
# sum of net parameters number
print("Number of trainable parameters in %s : %f" % (model_name, sum(p.numel() for p in net.parameters() if p.requires_grad)))
# print model structure
if only_print:
print(net)
return
# resume training
param_path = "./model/%s_%s_parameter.pt" % (model_name, "define" if self_define else "official")
if resume:
if os.path.exists(param_path):
net.load_state_dict(torch.load(param_path))
net.train()
print("Resume training " + model_name)
else:
print("Train %s from scratch" % model_name)
else:
print("Train %s from scratch" % model_name)
# define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# train on GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('train on %s' % device)
net.to(device)
running_loss = 0.0
train_losses = []
valid_losses = []
mini_batches = 125 * 5
for epoch in range(epochs):
for i, data in enumerate(trainloader, 0):
# get one batch
# inputs, labels = data
inputs, labels = data[0].to(device), data[1].to(device)
# switch model to training mode, clear gradient accumulators
net.train()
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % mini_batches == mini_batches - 1: # print and valid every <mini_batches> mini-batches
# validate model in validation dataset
valid_loss = valid(net, validloader, criterion, device)
print('[%d, %5d] train loss: %.3f, validset loss: %.3f' % (
epoch + 1, i + 1, running_loss / mini_batches, valid_loss))
train_losses.append(running_loss / mini_batches)
valid_losses.append(valid_loss)
running_loss = 0.0
# save parameters
torch.save(net.state_dict(), param_path)
# # save checkpoint
# torch.save({
# 'epoch': epoch,
# 'model_state_dict': net.state_dict(),
# 'optimizer_state_dict': optimizer.state_dict(),
# 'loss': loss
# }, "./checkpoints/epoch_" + str(epoch) + ".tar")
print('Finished Training, %d images in all' % (len(train_losses) * batch_size * mini_batches / epochs))
# draw loss curve
assert len(train_losses) == len(valid_losses)
loss_x = range(0, len(train_losses))
plt.plot(loss_x, train_losses, label="train loss")
plt.plot(loss_x, valid_losses, label="valid loss")
plt.title("Loss for every %d mini-batch" % mini_batches)
plt.xlabel("%d mini-batches" % mini_batches)
plt.ylabel("Loss")
plt.legend()
plt.savefig(model_name + "_loss.png")
plt.show()
def main():
parser = argparse.ArgumentParser()
mode_group = parser.add_mutually_exclusive_group(required=True)
mode_group.add_argument("--train",
action="store_true",
help="To train the network.")
mode_group.add_argument("--test",
action="store_true",
help="To test the network.")
mode_group.add_argument("--examples",
action="store_true",
help="To example MNIST data.")
parser.add_argument("--epochs",
default=10,
type=int,
help="Desired number of epochs.")
parser.add_argument("--dropout",
action="store_true",
help="Whether to use dropout or not.")
parser.add_argument("--uncertainty",
action="store_true",
help="Use uncertainty or not.")
uncertainty_type_group = parser.add_mutually_exclusive_group()
uncertainty_type_group.add_argument(
"--mse",
action="store_true",
help=
"Set this argument when using uncertainty. Sets loss function to Expected Mean Square Error."
)
uncertainty_type_group.add_argument(
"--digamma",
action="store_true",
help=
"Set this argument when using uncertainty. Sets loss function to Expected Cross Entropy."
)
uncertainty_type_group.add_argument(
"--log",
action="store_true",
help=
"Set this argument when using uncertainty. Sets loss function to Negative Log of the Expected Likelihood."
)
args = parser.parse_args()
if args.examples:
examples = enumerate(dataloaders["val"])
batch_idx, (example_data, example_targets) = next(examples)
fig = plt.figure()
for i in range(6):
plt.subplot(2, 3, i + 1)
plt.tight_layout()
plt.imshow(example_data[i][0], cmap="gray", interpolation="none")
plt.title("Ground Truth: {}".format(example_targets[i]))
plt.xticks([])
plt.yticks([])
plt.savefig("./images/examples.jpg")
elif args.train:
num_epochs = args.epochs
use_uncertainty = args.uncertainty
num_classes = 10
model = LeNet(dropout=args.dropout)
if use_uncertainty:
if args.digamma:
criterion = edl_digamma_loss
elif args.log:
criterion = edl_log_loss
elif args.mse:
criterion = edl_mse_loss
else:
parser.error(
"--uncertainty requires --mse, --log or --digamma.")
else:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.005)
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=7,
gamma=0.1)
device = get_device()
model = model.to(device)
model, metrics = train_model(model,
dataloaders,
num_classes,
criterion,
optimizer,
scheduler=exp_lr_scheduler,
num_epochs=num_epochs,
device=device,
uncertainty=use_uncertainty)
state = {
"epoch": num_epochs,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
}
if use_uncertainty:
if args.digamma:
torch.save(state, "./results/model_uncertainty_digamma.pt")
print("Saved: ./results/model_uncertainty_digamma.pt")
if args.log:
torch.save(state, "./results/model_uncertainty_log.pt")
print("Saved: ./results/model_uncertainty_log.pt")
if args.mse:
torch.save(state, "./results/model_uncertainty_mse.pt")
print("Saved: ./results/model_uncertainty_mse.pt")
else:
torch.save(state, "./results/model.pt")
print("Saved: ./results/model.pt")
elif args.test:
use_uncertainty = args.uncertainty
device = get_device()
model = LeNet()
model = model.to(device)
optimizer = optim.Adam(model.parameters())
if use_uncertainty:
if args.digamma:
checkpoint = torch.load(
"./results/model_uncertainty_digamma.pt")
filename = "./results/rotate_uncertainty_digamma.jpg"
if args.log:
checkpoint = torch.load("./results/model_uncertainty_log.pt")
filename = "./results/rotate_uncertainty_log.jpg"
if args.mse:
checkpoint = torch.load("./results/model_uncertainty_mse.pt")
filename = "./results/rotate_uncertainty_mse.jpg"
else:
checkpoint = torch.load("./results/model.pt")
filename = "./results/rotate.jpg"
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.eval()
rotating_image_classification(model,
digit_one,
filename,
uncertainty=use_uncertainty)
img = Image.open("./data/one.jpg").convert('L')
test_single_image(model, img, uncertainty=use_uncertainty)
def evalidation(model_name, testloader, classes, input_channel=3, self_define=True):
dataiter = iter(testloader)
images, labels = dataiter.next()
# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(batch_size)))
# load model parameter
assert model_name in ["LeNet", "VGG16", "ResNet", "DenseNet"]
param_path = "./model/%s_%s_parameter.pt" % (model_name, "define" if self_define else "official")
print("load model parameter from %s" % param_path)
if self_define:
if model_name == "LeNet":
net = LeNet(input_channel)
elif model_name == "VGG16":
net = VGG16(input_channel)
elif model_name == "ResNet":
net = ResNet(input_channel)
elif model_name == "DenseNet":
net = DenseNet(input_channel)
else:
if model_name == "LeNet":
net = LeNet(input_channel)
elif model_name == "VGG16":
net = models.vgg16_bn(pretrained=False, num_classes=10)
elif model_name == "ResNet":
net = models.resnet50(pretrained=False, num_classes=10)
elif model_name == "DenseNet":
net = models.DenseNet(num_classes=10)
net.load_state_dict(torch.load(param_path))
net.eval()
# predict
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(batch_size)))
# to gpu
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
# evaluate
class_correct = np.zeros(10)
class_total = np.zeros(10)
with torch.no_grad():
for data in testloader:
inputs, labels = data[0].to(device), data[1].to(device)
outputs = net(inputs)
_, predicted = torch.max(outputs, 1)
for i in range(batch_size):
label = labels[i]
class_total[label] += 1
if predicted[i] == label:
class_correct[label] += 1
print("\nEvery class precious: \n ",
' '.join("%5s : %2d %%\n" % (classes[i], 100 * class_correct[i]/class_total[i]) for i in range(len(classes))))
print("\n%d images in all, Total precious: %2d %%"
% (np.sum(class_total), 100 * np.sum(class_correct) / np.sum(class_total)))
torch.cuda.manual_seed(0)
torch.backends.cudnn.deterministic = True
MNIST_train = torchvision.datasets.MNIST('./', download=True, train=True)
MNIST_test = torchvision.datasets.MNIST('./', download=True, train=False)
X_train, y_train = MNIST_train.data, MNIST_train.targets
X_test, y_test = MNIST_test.data, MNIST_test.targets
X_train, X_test = X_train.float(), X_test.float()
X_train, X_test = X_train.unsqueeze(1), X_test.unsqueeze(1)
model = LeNet()
device = torch.device('cuda:0') if torch.cuda.is_available() else 'cpu'
model = model.to(device)
loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1.0e-3)
batch_size = 256
test_accuracy_history = []
test_loss_history = []
X_test = X_test.to(device)
y_test = y_test.to(device)
for epoch in range(1000):
order = np.random.permutation(len(X_train))
