def main():
# args & device
args = config.get_args()
if torch.cuda.is_available():
print('Train on GPU!')
device = torch.device("cuda")
else:
device = torch.device("cpu")
# dataset
assert args.dataset in ['cifar10', 'imagenet']
train_transform, valid_transform = data_transforms(args)
if args.dataset == 'cifar10':
trainset = torchvision.datasets.CIFAR10(root=os.path.join(args.data_dir, 'cifar'), train=True,
download=True, transform=train_transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,
shuffle=True, pin_memory=True, num_workers=8)
valset = torchvision.datasets.CIFAR10(root=os.path.join(args.data_dir, 'cifar'), train=False,
download=True, transform=valid_transform)
val_loader = torch.utils.data.DataLoader(valset, batch_size=args.batch_size,
shuffle=False, pin_memory=True, num_workers=8)
elif args.dataset == 'imagenet':
train_data_set = datasets.ImageNet(os.path.join(args.data_dir, 'ILSVRC2012', 'train'), train_transform)
val_data_set = datasets.ImageNet(os.path.join(args.data_dir, 'ILSVRC2012', 'valid'), valid_transform)
train_loader = torch.utils.data.DataLoader(train_data_set, batch_size=args.batch_size, shuffle=True,
num_workers=8, pin_memory=True, sampler=None)
val_loader = torch.utils.data.DataLoader(val_data_set, batch_size=args.batch_size, shuffle=False,
num_workers=8, pin_memory=True)
# SinglePath_OneShot
choice = [3, 1, 2, 1, 0, 1, 3, 3, 1, 3, 0, 1, 0, 3, 3, 3, 3, 3, 0, 3]
#[2, 0, 2, 3, 2, 2, 3, 1, 2, 1, 0, 1, 0, 3, 1, 0, 0, 2, 3, 2]
model = SinglePath_Network(args.dataset, args.resize, args.classes, args.layers, choice)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, args.momentum, args.weight_decay)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda epoch: 1 - (epoch / args.epochs))
# flops & params & structure
flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32),) if args.dataset == 'cifar10'
else (torch.randn(1, 3, 224, 224),), verbose=False)
# print(model)
print('Random Path of the Supernet: Params: %.2fM, Flops:%.2fM' % ((params / 1e6), (flops / 1e6)))
model = model.to(device)
summary(model, (3, 32, 32) if args.dataset == 'cifar10' else (3, 224, 224))
# train supernet
start = time.time()
for epoch in range(args.epochs):
train(args, epoch, train_loader, device, model, criterion, optimizer, scheduler, supernet=False)
scheduler.step()
if (epoch + 1) % args.val_interval == 0:
validate(args, epoch, val_loader, device, model, criterion, supernet=False)
utils.save_checkpoint({'state_dict': model.state_dict(), }, epoch + 1, tag=args.exp_name)
utils.time_record(start)
def run(name):
global x_train
global y_train
global x_dev
global y_dev
x_train, y_train = shuffle(x_train, y_train, random_state=1)
x_dev, y_dev = shuffle(x_dev, y_dev, random_state=1)
weights = model.train(x_train, y_train)
print("validating: " + name)
model.validate(weights, x_dev, y_dev)
print("writing test predictions for " + name)
results = model.get_test_results(x_test, weights)
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
pred_path = OUTPUT_DIR + "/test." + name + ".pred"
utils.print_results(results, pred_path)
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=8)
# random search
start = time.time()
best_acc = 0.0
acc_list = list()
best_choice = list()
for epoch in range(args.random_search):
choice = utils.random_choice(args.num_choices, args.layers)
top1_acc = validate(args,
epoch,
val_loader,
device,
model,
criterion,
super=True,
choice=choice)
acc_list.append(top1_acc)
if best_acc < top1_acc:
best_acc = top1_acc
best_choice = choice
print('acc_list:')
for i in acc_list:
print(i)
print('best_acc:{} \nbest_choice:{}'.format(best_acc, best_choice))
utils.plot_hist(acc_list, name=args.exp_name)
utils.time_record(start)
# Validation data
test_X = X[-val_size:]
test_y = y[-val_size:]
print(X[0].shape)
print(len(labels), "Classes")
net = model.Net(len(labels)).to(model.device)
print(net)
optimizer = optim.Adam(net.parameters(), lr=0.001)
loss_function = nn.MSELoss()
#Train model
model.train(net, BATCH_SIZE, EPOCHS, train_X, train_y, loss_function,
optimizer, True, test_X, test_y)
accuracy, validation_loss = model.validate(net, test_X, test_y,
loss_function,
BATCH_SIZE) #Validation
torch.save(net.state_dict(), 'mytraining.pt') #Save model
print(f'Validation Accuracy {accuracy}, Validation Loss {validation_loss}'
) #Accuracy
"""
cv2.imshow(f'{training_data[0][1]}', training_data[0][0])
cv2.waitKey(0)
cv2.destroyAllWindows()
"""
if TEST_DATA:
test_data = np.load(f'{data.SketchData.TEST_PATH}/test_data.npy',
allow_pickle=True)
labels = np.load(f'{data.SketchData.TRAIN_PATH}/labels.npy',
allow_pickle=True)