def main(num_epochs=200,
learning_rate=0.005,
momentum=0.5,
log_interval=500,
*args,
**kwargs):
train_loader, test_loader = loaders.loader(batch_size_train=100,
batch_size_test=1000)
# Train the model
total_step = len(train_loader)
curr_lr1 = learning_rate
model1 = VGG().to(device)
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer1 = torch.optim.Adam(model1.parameters(), lr=learning_rate)
# Train the model
total_step = len(train_loader)
best_accuracy1 = 0
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
# Forward
outputs = model1(images)
loss1 = criterion(outputs, labels)
# Backward and optimize
optimizer1.zero_grad()
loss1.backward()
optimizer1.step()
if i == 499:
print(
"Ordinary Epoch [{}/{}], Step [{}/{}] Loss: {:.4f}".format(
epoch + 1, num_epochs, i + 1, total_step,
loss1.item()))
# Test the model
model1.eval()
with torch.no_grad():
correct1 = 0
total1 = 0
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model1(images)
_, predicted = torch.max(outputs.data, 1)
total1 += labels.size(0)
correct1 += (predicted == labels).sum().item()
if best_accuracy1 >= correct1 / total1:
curr_lr1 = learning_rate * np.asscalar(
pow(np.random.rand(1), 3))
update_lr(optimizer1, curr_lr1)
print('Test Accuracy of NN: {} % Best: {} %'.format(
100 * correct1 / total1, 100 * best_accuracy1))
else:
best_accuracy1 = correct1 / total1
net_opt1 = model1
print('Test Accuracy of NN: {} % (improvement)'.format(
100 * correct1 / total1))
model1.train()
model_id = 0 # Change this to correspond to the model in the list
if model_id == 0:
model = VGG3D()
elif model_id == 1:
model = resnet34()
elif model_id == 2:
model = ResNet2D(3)
elif model_id == 3:
model = VGG(3)
train_size = 15
test_size = 15
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.train()
model = model.to(device)
model = torch.load("model_m.pt")
model.eval()
train_loader = get_data_loader(train = True, batch_size = train_size,
split = 'train', model = models[model_id])
test_loader = get_data_loader(train = False, batch_size = test_size,
split = 'test', model = models[model_id])
_ , train_acc = test(model, train_loader, device)
print("Final Train Accuracy: ", train_acc)
_ , test_acc = test(model, test_loader, device)
print("Final Accuracy: ", test_acc)
expression_queue = Queue()
expression_faceID_queue = Queue()
# ================= DEFINITION =====================
similarityThreshold = 0.7
currentFaceID = []
faceIdToName = {}
faceLib = {}
expressionDict = {}
# ===================== LOAD FER MODEL =================
net = VGG('VGG19')
checkpoint = torch.load('FER2013_VGG19/expression_recognition_model.t7')
net.load_state_dict(checkpoint['net'])
net.cuda()
net.eval()
# =================== RECORD START TIME ================
stat_time = time.time()
# =================== FER INPUT_SIZE ADJUSTMENT ===========
cut_size = 44
transform_test = transforms.Compose([
transforms.TenCrop(cut_size),
transforms.Lambda(lambda crops: torch.stack(
[transforms.ToTensor()(crop) for crop in crops])),
])
# ================== EMOTION CATEGORY ================
EMOTIONS = ['生气', '厌恶', '害怕', '开心', '难过', '惊讶', '平静']
def main():
"""
This code sets up the data and loads the model obtained after ADMM based training
for retraining to enforce pruning constraints.
The function retrains_model present in the utils file enforces the hard sparsity
constraints on the weights on the model obtained after ADMM based training
and then retrains the model to improve the accuracy.
"""
#model = LeNet5()
model = VGG(n_class=10)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
Path = 'saved_model/admm_model/cifar10_vgg_acc_0.688' # Path to the saved model after ADMM based training
model.load_state_dict(torch.load(Path))
model.to(device)
#data_transforms = transforms.Compose([transforms.CenterCrop(32),transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])
train_data = datasets.CIFAR10('data/', train=True, download=False,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
]))
#train_data = datasets.MNIST(root='data/',download=False,train=True,transform=data_transforms)
# Splitting the training dataset into training and validation dataset
N_train = len(train_data)
val_split = 0.1
N_val = int(val_split*N_train)
train_data,val_data = torch.utils.data.random_split(train_data,(N_train-N_val,N_val))
## Test data
test_data = datasets.CIFAR10('data/', train=False, download=False,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
]))
#test_data = datasets.MNIST(root='data/',download=False,train=False,transform=data_transforms)
batch_size = 128
num_epochs = 20
log_step = 100
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=1e-2)
#optimizer = torch.optim.SGD(model.parameters(), lr =5e-4,momentum =0.9, weight_decay = 5e-4 )
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,milestones = [10], gamma = 0.1)
####### Re-Training ##############
# Parameters
prune_type = 'filter'
# Number of non-zero filters at each convolutional layer
l = {'conv1':32,'conv2':64,'conv3':128,'conv4':128,'conv5':256,'conv6':256,'conv7':256,'conv8':256}
retrain_model(model,train_data,val_data,batch_size,loss_fn,num_epochs,log_step,optimizer,scheduler,l,prune_type,device)
# Check the test accuracy
model.eval()
test_accuracy = eval_accuracy_data(test_data,model,batch_size,device)
print('Test accuracy is',test_accuracy)
def main():
"""
This code implements the ADMM based training of a CNN.
"""
#model = LeNet5()
model = VGG(n_class=10)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
Path = 'saved_model/pre_train_models/cifar10_vgg_acc_0.943' # Path to the baseline model
model.load_state_dict(torch.load(Path))
model.to(device)
#data_transforms = transforms.Compose([transforms.CenterCrop(32),transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])
train_data = datasets.CIFAR10('data/',
train=True,
download=False,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
]))
#train_data = datasets.MNIST(root='data/',download=False,train=True,transform=data_transforms)
"""
N_train = len(train_data)
val_split = 0.1
N_val = int(val_split*N_train)
train_data,val_data = torch.utils.data.random_split(train_data,(N_train-N_val,N_val))
"""
## Test data
test_data = datasets.CIFAR10('data/',
train=False,
download=False,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
]))
#test_data = datasets.MNIST(root='data/',download=False,train=False,transform=data_transforms)
batch_size = 128
num_epochs = 50
log_step = 100
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
#optimizer = torch.optim.SGD(model.parameters(), lr =5e-4,momentum =0.9, weight_decay = 5e-4 )
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[15, 30],
gamma=0.1)
####### ADMM Training ##############
## Parameters
fc_prune = False # True if the fully connected layers are also pruned
prune_type = 'filter' # Type of structural pruning at the convolutional layers
# Number of non zero filters at each convolutional layer
l = {
'conv1': 32,
'conv2': 64,
'conv3': 128,
'conv4': 128,
'conv5': 256,
'conv6': 256,
'conv7': 256,
'conv8': 256
}
# ADMM parameters
rho_val = 1.5e-3
num_admm_steps = 10
Z = {}
U = {}
rho = {}
best_accuracy = 0
all_acc = False
## Initialization of the variable Z and dual variable U
for name_net in model.named_modules():
name, net = name_net
if isinstance(net, nn.Conv2d):
Z[name] = net.weight.clone().detach().requires_grad_(False)
Z[name] = Projection_structured(Z[name], l[name], prune_type)
U[name] = torch.zeros_like(net.weight, requires_grad=False)
rho[name] = rho_val
elif fc_prune and isinstance(net, nn.Linear):
Z[name] = net.weight.clone().detach().requires_grad_(False)
l_unst = int(len(net.weight.data.reshape(-1, )) * prune_ratio)
Z[name], _ = Projection_unstructured(Z[name], l_unst)
U[name] = torch.zeros_like(net.weight, requires_grad=False)
## ADMM loop
for i in range(num_admm_steps):
print('ADMM step number {}'.format(i))
# First train the VGG model
train_model_admm(model, train_data, batch_size, loss_fn, optimizer,
scheduler, num_epochs, log_step, Z, U, rho, fc_prune,
device)
# Update the variable Z
for name_net in model.named_modules():
name, net = name_net
if isinstance(net, nn.Conv2d):
Z[name] = Projection_structured(net.weight.detach() + U[name],
l[name], prune_type)
elif fc_prune and isinstance(net, nn.Linear):
l_unst = int(len(net.weight.data.reshape(-1, )) * prune_ratio)
Z[name], _ = Projection_unstructured(
net.weight.detach() + U[name], l_unst)
# Updating the dual variable U
for name_net in model.named_modules():
name, net = name_net
if isinstance(net, nn.Conv2d):
U[name] = U[name] + net.weight.detach() - Z[name]
elif fc_prune and isinstance(net, nn.Linear):
U[name] = U[name] + net.weight.detach() - Z[name]
## Check the test accuracy
model.eval()
test_accuracy = eval_accuracy_data(test_data, model, batch_size,
device)
print('Test accuracy is', test_accuracy)
if test_accuracy > best_accuracy:
print(
'Saving model with test accuracy {:.3f}'.format(test_accuracy))
torch.save(
model.state_dict(),
'saved_model/admm_model/cifar10_vgg_acc_{:.3f}'.format(
test_accuracy))
if all_acc:
print('Removing model with test accuracy {:.3f}'.format(
best_accuracy))
os.remove(
'saved_model/admm_model/cifar10_vgg_acc_{:.3f}'.format(
best_accuracy))
best_accuracy = test_accuracy
all_acc = True