def init_local_model(self, device):
### Select the model
if self.params.t_model_version == 'resnet18':
self.m_local = (resnet.ResNet18()).to(device) if self.params.t_cuda else (resnet.ResNet18())
elif self.params.t_model_version == 'alexnet':
self.m_local = (alexnet.AlexNet()).to(device) if self.params.t_cuda else (alexnet.AlexNet())
elif self.params.t_model_version == 'cnn':
self.m_local = (cnn.CNN()).to(device) if self.params.t_cuda else (cnn.CNN())
if args.quantization == 'bwn':
model = nin.BWNNet()
model_name = args.arch + '_bwn.pth.tar'
f = open(os.path.join('log', 'log_' + args.arch + "_bwn"), 'w')
elif args.quantization == 'xnor':
model = nin.XNORNet()
model_name = args.arch + '_xnor.pth.tar'
f = open(os.path.join('log', 'log_' + args.arch + "_xnor"), 'w')
elif args.quantization == 'joint_bwn':
if args.arch == 'nin':
model = nin.BWNNet()
elif args.arch == 'alex':
model = alexnet.AlexNet()
name = args.arch + "_joint_bwn_" + str(args.slice)
if args.binarize_first_layer == True:
name = name + "_bf"
if args.binarize_last_layer == True:
name = name + "_bl"
model_name = name + '.pth.tar'
f = open(os.path.join('log', 'log_' + name), 'w')
elif args.quantization == 'joint_xnor':
model = nin.XNORNet()
name = args.arch + "_joint_xnor_" + str(args.slice)
if args.binarize_first_layer == True:
name = name + "_bf"
if args.binarize_last_layer == True:
def train_baseline_KD(self):
if self.params.model_version == 'resnet18':
model = resnet.ResNet18().cuda(
) if self.params.cuda else resnet.ResNet18()
elif self.params.model_version == 'alexnet':
model = alexnet.AlexNet().cuda(
) if self.params.cuda else alexnet.AlexNet()
optimizer = optim.SGD(model.parameters(),
lr=self.params.learning_rate,
momentum=0.9,
weight_decay=5e-4)
loss_function_KD = utils.loss_function_kd
metrics = utils.metrics
teacher_model = resnet.ResNet18()
teacher_checkpoint = 'experiments/baseline_standalone_resnet18/best.pth.tar'
teacher_model = teacher_model.cuda(
) if self.params.cuda else teacher_model
utils.load_checkpoint(teacher_checkpoint, teacher_model)
# Train the model with KD
logging.info("Experiment - model version: {}".format(
self.params.model_version))
logging.info("Starting training for {} epoch(s)".format(
self.params.num_epochs))
logging.info(
"First, loading the teacher model and computing its outputs...")
best_valid_acc = 0.0
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=150,
gamma=0.1)
for epoch in range(self.params.num_epochs):
logging.info("Epoch {}/{}".format(epoch + 1,
self.params.num_epochs))
train_metrics = baseline_KD.train_kd(model, teacher_model,
optimizer, loss_function_KD,
self.trainloader, metrics,
self.params)
scheduler.step()
valid_metrics = baseline_KD.evaluate_kd(model, teacher_model,
loss_function_KD,
self.testloader, metrics,
self.params)
valid_acc = valid_metrics['accuracy']
is_best = valid_acc >= best_valid_acc
# Record experiment results
with open(self.model_dir + "/result.csv", 'a') as f:
writer = csv.writer(f)
row = [
train_metrics['loss'], train_metrics['accuracy'],
valid_metrics['loss'], valid_metrics['accuracy']
]
writer.writerow(row)
# Save weights
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict()
},
is_best=is_best,
checkpoint=self.model_dir)
# If best_eval, best_save_path
if is_best:
logging.info("- Found new best accuracy")
best_valid_acc = valid_acc
# Save best valid metrics in a JSON file in the model directory
best_json_path = os.path.join(
self.model_dir, "metrics_valid_best_weights.json")
utils.save_dict_to_json(valid_metrics, best_json_path)
# Save latest valid metrics in a JSON file in the model directory
last_json_path = os.path.join(self.model_dir,
"metrics_valid_last_weights.json")
utils.save_dict_to_json(valid_metrics, last_json_path)
def train_baseline_standalone(self):
if self.params.model_version == 'resnet18':
model = resnet.ResNet18().cuda(
) if self.params.cuda else resnet.ResNet18()
elif self.params.model_version == 'alexnet':
model = alexnet.AlexNet().cuda(
) if self.params.cuda else alexnet.AlexNet()
optimizer = optim.SGD(model.parameters(),
lr=self.params.learning_rate,
momentum=0.9,
weight_decay=5e-4)
loss_function = utils.loss_function
metrics = utils.metrics
# Train the model
logging.info("Starting training for {} epoch(s)".format(
self.params.num_epochs))
best_valid_acc = 0
if self.params.model_version == "resnet18":
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=150,
gamma=0.1)
elif self.params.model_version == "alexnet":
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=100,
gamma=0.2)
for epoch in range(self.params.num_epochs):
logging.info("Epoch {}/{}".format(epoch + 1,
self.params.num_epochs))
train_metrics = baseline_standalone.train(model, optimizer,
loss_function,
self.trainloader,
metrics, self.params)
scheduler.step()
valid_metrics = baseline_standalone.evaluate(
model, loss_function, self.testloader, metrics, self.params)
valid_acc = valid_metrics['accuracy']
is_best = valid_acc >= best_valid_acc
# Record experiment results
with open(self.model_dir + "/result.csv", 'a') as f:
writer = csv.writer(f)
row = [
train_metrics['loss'], train_metrics['accuracy'],
valid_metrics['loss'], valid_metrics['accuracy']
]
writer.writerow(row)
# Save weights
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict()
},
is_best=is_best,
checkpoint=self.model_dir)
if is_best:
logging.info("- Found new best accuracy")
best_valid_acc = valid_acc
# Save best validation metrics in a json file in the model directory
best_json_path = os.path.join(
self.model_dir, "metrics_valid_best_weights.json")
utils.save_dict_to_json(valid_metrics, best_json_path)
# Save latest valid metrics in a json file in the model directory
last_json_path = os.path.join(self.model_dir,
"metrics_valid_last_weights.json")
utils.save_dict_to_json(valid_metrics, last_json_path)
def main(args):
BATCH_SIZE = args.batch_size
LR = args.learning_rate
EPOCH = args.epoch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
use_gpu = torch.cuda.is_available()
data_transforms = {
transforms.Compose([
transforms.Resize(320),
transforms.CenterCrop(299),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
}
transform = transforms.Compose([
transforms.Resize(size=(227, 227)),
transforms.RandomRotation(20),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), # 将图片转换为Tensor,归一化至[0,1]
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.ImageFolder(root=args.train_images,
transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True)
# 从文件夹中读取validation数据
validation_dataset = torchvision.datasets.ImageFolder(
root=args.test_images, transform=transform)
print(validation_dataset.class_to_idx)
test_loader = torch.utils.data.DataLoader(validation_dataset,
batch_size=BATCH_SIZE,
shuffle=True)
if args.model_name == "densenet":
Net = densenet.DenseNet().to(device)
if args.model_name == "alexnet":
Net = alexnet.AlexNet().to(device)
if args.model_name == "googlenet":
Net = googlenet.GoogLeNet().to(device)
if args.model_name == "mobilenet":
Net = mobilenet.MobileNetV2().to(device)
if args.model_name == "mnasnet":
Net = mnasnet.mnasnet1_0().to(device)
if args.model_name == "squeezenet":
Net = squeezenet.SqueezeNet().to(device)
if args.model_name == "resnet":
Net = resnet.resnet50().to(device)
if args.model_name == "vgg":
Net = vgg.vgg19().to(device)
if args.model_name == "shufflenetv2":
Net = shufflenetv2.shufflenet_v2_x1_0().to(device)
criterion = nn.CrossEntropyLoss()
opti = torch.optim.Adam(Net.parameters(), lr=LR)
if __name__ == '__main__':
Accuracy_list = []
Loss_list = []
for epoch in range(EPOCH):
sum_loss = 0.0
correct1 = 0
total1 = 0
for i, (images, labels) in enumerate(train_loader):
num_images = images.size(0)
images = Variable(images.to(device))
labels = Variable(labels.to(device))
if args.model_name == 'googlenet':
out = Net(images)
out = out[0]
else:
out = Net(images)
_, predicted = torch.max(out.data, 1)
total1 += labels.size(0)
correct1 += (predicted == labels).sum().item()
loss = criterion(out, labels)
print(loss)
opti.zero_grad()
loss.backward()
opti.step()
# 每训练100个batch打印一次平均loss
sum_loss += loss.item()
if i % 10 == 9:
print('train loss [%d, %d] loss: %.03f' %
(epoch + 1, i + 1, sum_loss / 2000))
print("train acc %.03f" % (100.0 * correct1 / total1))
sum_loss = 0.0
Accuracy_list.append(100.0 * correct1 / total1)
print('accurary={}'.format(100.0 * correct1 / total1))
Loss_list.append(loss.item())
x1 = range(0, EPOCH)
x2 = range(0, EPOCH)
y1 = Accuracy_list
y2 = Loss_list
total_test = 0
correct_test = 0
for i, (images, labels) in enumerate(test_loader):
start_time = time.time()
print('time_start', start_time)
num_images = images.size(0)
print('num_images', num_images)
images = Variable(images.to(device))
labels = Variable(labels.to(device))
print("GroundTruth", labels)
if args.model_name == 'googlenet':
out = Net(images)[0]
out = out[0]
else:
out = Net(images)
_, predicted = torch.max(out.data, 1)
print("predicted", predicted)
correct_test += (predicted == labels).sum().item()
total_test += labels.size(0)
print('time_usage', (time.time() - start_time) / args.batch_size)
print('total_test', total_test)
print('correct_test', correct_test)
print('accurary={}'.format(100.0 * correct_test / total_test))
plt.subplot(2, 1, 1)
plt.plot(x1, y1, 'o-')
plt.title('Train accuracy vs. epoches')
plt.ylabel('Train accuracy')
plt.subplot(2, 1, 2)
plt.plot(x2, y2, '.-')
plt.xlabel('Train loss vs. epoches')
plt.ylabel('Train loss')
# plt.savefig("accuracy_epoch" + str(EPOCH) + ".png")
plt.savefig(args.output_dir + '/' + 'accuracy_epoch' + str(EPOCH) +
'.png')
plt.show()
torch.save(args.output_dir, args.model_name + '.pth')
from models import preact_resnet_model
from models import alexnet
from models import lenet
_MODEL_DICT = {
'resnet18': resnet_model.ResNet18(),
'resnet34': resnet_model.ResNet34(),
'resnet50': resnet_model.ResNet50(),
'resnet101': resnet_model.ResNet101(),
'resnet152': resnet_model.ResNet152(),
'preact_resnet18': preact_resnet_model.PreActResNet18(),
'preact_resnet34': preact_resnet_model.PreActResNet34(),
'preact_resnet50': preact_resnet_model.PreActResNet50(),
'preact_resnet101': preact_resnet_model.PreActResNet101(),
'preact_resnet152': preact_resnet_model.PreActResNet152(),
'alexnet': alexnet.AlexNet(),
'lenet': lenet.LeNet()
}
def _invalid_model_name():
raise ValueError("Not a valid model name")
def fetch_teacher(model_name, model_dicts=_MODEL_DICT):
return model_dicts[model_name]
def fetch_student(model_name, model_dicts=_MODEL_DICT):
return model_dicts[model_name]
parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
parser.add_argument('--threshold',
default=0.7,
type=float,
help='learning rate')
parser.add_argument('--input',
default='test_image/3_1.jpg',
help='Input Image')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
args = parser.parse_args()
classes = ['Red', 'Yellow', 'Green', 'Left', 'Right', 'Straight']
img_size = 224
print('==> Building model..')
net = alexnet.AlexNet(num_classes=len(classes))
#net = mobilenet.MobileNet(num_classes=4)
net = net.to(device)
net.eval()
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
#Load checkpoint
print('==> Resuming from checkpoint..')
assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
checkpoint = torch.load('./checkpoint/ckpt_old_i.t7')
net.load_state_dict(checkpoint['net'])
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch']
def loadCoeffIdx(pathName, modelName):
skipLinear = False
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if modelName == "vgg16_c10":
model = vgg16.VGG16()
elif modelName == "vgg16_c10_old":
model = vgg16_old.VGG16()
elif modelName == "resnet50_c10":
model = resnet.ResNet50()
elif modelName == "inceptionv3_c10":
model = inception_v3_c10.inception_v3()
elif modelName == "alexnet_in":
#model = torch.hub.load('pytorch/vision', 'alexnet', pretrained=True)
model = alexnet.AlexNet()
skipLinear = True
elif modelName == "vgg16_in":
model = vgg_in.vgg16()
elif modelName == "resnet50_in":
model = resnet_in.resnet50()
elif modelName == "transformer_wmt":
translator = transformer_translate.create_model(
) # returns Transformer()
translator = translator.to(device)
model = translator.model
else:
print("Model not supported")
exit(1)
model = model.to(device)
if ("vgg16" in modelName or "alexnet" in modelName):
model = replace_vgg16(model, skipLinear)
elif "transformer" in modelName:
#transformer_train.replace_with_pruned(model, "model", prune_attention=True, prune_only_attention=False)
pass
else:
replace_with_pruned(model, "model", skipLinear)
if "_in" in modelName:
if not torch.distributed.is_initialized():
port = np.random.randint(10000, 65536)
torch.distributed.init_process_group(
backend='nccl',
init_method='tcp://127.0.0.1:%d' % port,
rank=0,
world_size=1)
model = torch.nn.parallel.DistributedDataParallel(model)
if not ("transformer" in modelName):
model.load_state_dict(torch.load(pathName))
model.eval()
#prune(model, method="cascade", q=1.0)
layers = []
widths = [3]
numConv = 0
assert isinstance(model, nn.Module)
for n, m in model.named_modules():
print(type(m))
if isinstance(m, torch.nn.Conv2d):
numConv += 1
if isinstance(m, torch.nn.Linear):
numConv += 1
if isinstance(m, PrunedConv):
#print(m.mask)
#layer = m.mask.view((m.out_channels, -1)).detach().cpu().numpy()
layer = m.conv.weight.view(
(m.out_channels, -1)).detach().cpu().numpy()
elif isinstance(m, PrunedLinear) or isinstance(
m, CSP.pruned_layers.PrunedLinear) and (not skipLinear):
#print(m.mask)
#layer = m.mask.view((m.out_features, -1)).detach().cpu().numpy()
layer = m.linear.weight.view(
(m.out_features, -1)).detach().cpu().numpy()
else:
continue
#print(n)
#print(layer.shape)
#layer = layer > 0
widths.append(len(layer))
layers.append(layer)
#layerMask = layer > 0
#layerMask = np.transpose(layerMask)
#print("NUM CONV: {}".format(numConv))
#print("NUM EXTRACTED LAYERS: {}".format(len(layers)))
if "transformer" in modelName:
for i in range(11):
for j in range(36, 97):
layers.append(layers[j])
return layers
def exportData(pathName, modelName):
global skipLinear
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if modelName == "vgg16_c10":
model = vgg16.VGG16()
elif modelName == "vgg16_c10_old":
model = vgg16_old.VGG16()
elif modelName == "resnet50_c10":
model = resnet.ResNet50()
elif modelName == "inceptionv3_c10":
model = inception_v3_c10.inception_v3()
elif modelName == "alexnet_in":
#model = torch.hub.load('pytorch/vision', 'alexnet', pretrained=True)
model = alexnet.AlexNet()
skipLinear = True
elif modelName == "vgg16_in":
model = vgg_in.vgg16()
skipLinear = True
elif modelName == "resnet50_in":
model = resnet_in.resnet50()
elif modelName == "transformer_wmt":
translator = transformer_translate.create_model(
) # returns Transformer()
translator = translator.to(device)
model = translator.model
else:
print("Model not supported")
exit(1)
model = model.to(device)
if ("vgg16" in modelName or "alexnet" in modelName):
model = replace_vgg16(model, skipLinear)
elif "transformer" in modelName:
#transformer_train.replace_with_pruned(model, "model", prune_attention=True, prune_only_attention=False)
pass
else:
replace_with_pruned(model, "model", skipLinear)
if "_in" in modelName:
if not torch.distributed.is_initialized():
port = np.random.randint(10000, 65536)
torch.distributed.init_process_group(
backend='nccl',
init_method='tcp://127.0.0.1:%d' % port,
rank=0,
world_size=1)
model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=[0],
output_device=0)
if not ("transformer" in modelName):
model.load_state_dict(torch.load(pathName))
model.eval()
#prune(model, method="cascade", q=1.0)
#layer_num = 0
#save_dir = "/root/hostCurUser/reproduce/DNNsim/net_traces/ResNet50_ImageNet_CSP/"
#lstModules = list( model.named_modules())
#for n, m in model.named_modules():
#for i in range(len(lstModules)):
# if isinstance(lstModules[i], nn.Conv2d) or isinstance(lstModules[i], nn.Linear):
# model[i] = DataExporter(lstModules[i], save_dir, layer_num)
# layer_num += 1
f = open(MODEL_PATH + "model.csv", "w")
fscale = open(SCALE_PATH, "w")
fscale.write(
"Layer name, IFMAP Height, IFMAP Width, Filter Height, Filter Width, Channels, Num Filter, Strides,\n"
)
layer_idx = 0
models = [] # for SparTen
layers = []
acts = []
weights = []
paddings = []
strides = []
idxs = []
def extract(module, input):
#if module.extracted:
# return
if len(input[0].shape) < 4:
if not ("transformer" in modelName):
return
try:
a = input[0].detach().cpu().reshape(1, module.in_features, 1,
1)
except:
a = input[0].detach().cpu().reshape(-1, 1, 1)
a = a[:module.in_features]
a = a.reshape(1, module.in_features, 1, 1)
else:
a = input[0].detach().cpu()
acts.append(a)
if isinstance(module, torch.nn.Conv2d):
layer = module.weight.view(
(module.out_channels, -1)).detach().cpu().numpy()
weight = module.weight.detach().cpu().numpy()
tp = "conv"
stride = str(max(module.stride[0], module.stride[1]))
padding = str(max(module.padding[0], module.padding[1]))
in_channels = module.in_channels
out_channels = module.out_channels
kernel_size = module.kernel_size
padding_st = module.padding
stride_st = module.stride
elif isinstance(module, torch.nn.Linear) and (not skipLinearExport):
layer = module.weight.view(
(module.out_features, -1)).detach().cpu().numpy()
weight = module.weight.detach().cpu().reshape(
module.weight.shape[0], module.weight.shape[1], 1, 1).numpy()
tp = "fc"
stride = str(1)
padding = str(0)
in_channels = module.in_features
out_channels = module.out_features
kernel_size = (1, 1)
padding_st = (0, 0)
stride_st = (1, 1)
else:
print("{} does not exist".format(module))
exit(1)
name = str(module.layer_idx)
weights.append(weight)
paddings.append(int(padding))
strides.append(int(stride))
f.write(name + "," + tp + "," + stride + "," + padding + ",\n")
layers.append(layer)
models.append({
'in_channels': in_channels,
'out_channels': out_channels,
'kernel': kernel_size,
'name': tp + name,
'padding': padding_st,
'weights': weight,
'IFM': a.cpu().numpy(),
'stride': stride_st
})
idxs.append(module.layer_idx)
#module.extracted = True
#replace_with_exporter(model, "model", f, fscale)
for n, m in model.named_modules():
if isinstance(m, torch.nn.Conv2d):
weight = m.weight.detach().cpu().numpy()
if weight.shape[2] != weight.shape[3]:
continue
#weights.append(weight)
#m.extracted = False
m.register_forward_pre_hook(extract)
#name = str(layer_idx)
#tp = "conv"
#stride = str(max(m.stride[0], m.stride[1]))
#padding = str(max(m.padding[0], m.padding[1]))
#paddings.append(int(padding))
#f.write(name+"," + tp+"," + stride+"," + padding + ",\n")
m.layer_idx = layer_idx
layer_idx += 1
elif isinstance(m, torch.nn.Linear):
if not ("transformer" in modelName):
continue
#weight = m.weight.detach().cpu().reshape(m.weight.shape[0], m.weight.shape[1], 1, 1).numpy()
#weights.append(weight)
#m.extracted = False
m.register_forward_pre_hook(extract)
#name = str(layer_idx)
#tp = "fc"
#stride = str(1)
#padding = str(0)
#paddings.append(int(padding))
#f.write(name+"," + tp+"," + stride+"," + padding + ",\n")
m.layer_idx = layer_idx
layer_idx += 1
if "_in" in modelName:
IFM = torch.rand(1, 3, 224, 224).cuda()
model(IFM)
elif "_c10" in modelName:
IFM = torch.rand(1, 3, 32, 32).cuda()
model(IFM)
elif "_wmt" in modelName:
src_seq = [4556, 4562, 4560, 4557, 4712, 1894, 15, 4564, 4620, 0, 5]
pred_seq = translator.translate_sentence(
torch.LongTensor([src_seq]).to(device))
print(pred_seq)
else:
print("Dataset not supported")
exit(1)
print(len(acts))
print(len(weights))
#layers = loadCoeffIdx(pathName, modelName)
with open(ST_PATH + modelName + ".h5", "wb") as f:
pickle.dump(models, f)
sq_ptrs = []
out_channels = []
for layer in layers:
sp, oc = squeezeCoeffIdxTotal(layer, len(layer))
out_channels.append(oc)
#"""
i = 0
for idx in range(len(acts)):
x_save = acts[idx]
weight_save = weights[idx]
np.save(EXPORT_PATH + "act-" + str(idx) + "-0.npy", x_save)
np.save(EXPORT_PATH + "wgt-" + str(idx) + ".npy", weight_save)
x_save[x_save == 0] = int(0)
x_save[x_save != 0] = int(1)
np.savetxt(CX_PATH + "act" + str(idx) + ".csv",
x_save.reshape(-1),
delimiter=",",
fmt="%d")
weight_save[weight_save == 0] = int(0)
weight_save[weight_save != 0] = int(1)
np.savetxt(CX_PATH + "Conv2D_" + str(idx) + ".csv",
weight_save.reshape(weight_save.shape[0], -1),
delimiter=",",
fmt="%d")
if x_save.shape[2] > 1:
name = "Conv" + str(idx) #str(idxs[idx])
IFM_height = str(x_save.shape[2] + (2 * paddings[idx]))
IFM_width = str(x_save.shape[3] + (2 * paddings[idx]))
filt_height = str(weight_save.shape[2])
filt_width = str(weight_save.shape[3])
else:
name = "FC" + str(idx) #str(idxs[idx])
IFM_height = str(1)
IFM_width = str(1)
filt_height = str(1)
filt_width = str(1)
channels = str(weight_save.shape[1])
if ("resnet50" in modelName) and (idx == 4 or idx == 15 or idx == 29
or idx == 49):
num_filt = str(weight_save.shape[0])
else:
num_filt = str(out_channels[i])
i += 1
fscale.write(name + ',\t' + IFM_height + ',\t' + IFM_width + ',\t' +
filt_height + ',\t' + filt_width + ',\t' + channels +
',\t' + num_filt + ',\t' + str(strides[idx]) + ',\n')
fscale.close()
f.close()
#"""
cxf = open(CX_PATH + "ModelShape.txt", "w")
cxf.write("LayerName\tLayerID\tInputShape\tOutputShape\tKernelShape\n")
cx_layers = []
layer_idx = 0
for n, m in model.named_modules():
if isinstance(m, torch.nn.Conv2d):
if m.weight.shape[2] != m.weight.shape[3]:
continue
curr = "Conv\t" + str(layer_idx) + "\t" + str(
tuple(acts[layer_idx].shape)).replace('(', '').replace(
')', '').replace(' ', '') + "\t" + str(
tuple(acts[layer_idx].shape)).replace('(', '').replace(
')', '').replace(' ', '') + "\t" + str(
tuple(
m.weight.shape)).replace('(', '').replace(
')', '').replace(' ', '') + "\n"
cxf.write(curr)
cx_layers.append(curr)
layer_idx += 1
if isinstance(m, torch.nn.Linear) and (not skipLinear):
curr = "FC\t" + str(idxs[layer_idx]) + "\t" + str(
tuple(acts[layer_idx].shape)).replace('(', '').replace(
')', '').replace(' ', '') + "\t" + str(
tuple(acts[layer_idx].shape)).replace('(', '').replace(
')', '').replace(' ', '') + "\t" + str(
tuple(
m.weight.shape)).replace('(', '').replace(
')', '').replace(' ', '') + ",1,1\n"
cxf.write(curr)
cx_layers.append(curr)
layer_idx += 1
if "transformer" in modelName:
for i in range(11):
for j in range(36, 97):
cxf.write(cx_layers[j])
cxf.close()
return
_, preds = outputs.max(1)
total += targets.size(0)
correct += preds.eq(targets).sum().item()
test_loss /= len(data_loader.dataset)
acc = 100. * correct / total
print("\n Test set: Average loss: {:.4f}, Accuracy: {}/{} {:.4f}\n".format(
test_loss, correct, total, correct * 1.0 / total))
if __name__ == '__main__':
args = set_args()
args.cuda = torch.cuda.is_available()
torch.manual_seed(args.seed)
model = alexnet.AlexNet(num_classes=len(cifar10_classes))
# model = vgg.VGG('VGG16', len(cifar10_classes))
# model = vgg.VGG('VGG19', len(cifar10_classes))
# model = resnet.ResNet50(num_classes=len(cifar10_classes))
# model = resnet.ResNet101(num_classes=len(cifar10_classes))
# model = densenet.DenseNet121(num_classes=len(cifar10_classes))
# model = densenet.DenseNet201(num_classes=len(cifar10_classes))
if args.cuda:
torch.cuda.manual_seed(args.seed)
model.cuda(args.device_id)
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
# dataloader
data_loader = test_loader(args)
train_loader = preprocess.train_loader
val_loader = preprocess.val_loader
train_len = preprocess.train_len
val_len = preprocess.val_len
writer = SummaryWriter(preprocess.save_mpdel_path)
#data_root = os.getcwd()
device = torch.device(
cfg_content['device'] if torch.cuda.is_available() else "cpu")
# load model######################### manuall modify###############################
#net = resnet.resnet34(cfg.num_classes)
#net = vgg.VGG(cfg_content['model']['name'], cfg_content['dataset']['n_classes'])
#net = lenet.LeNet(cfg_content['dataset']['n_classes'])
net = alexnet.AlexNet(cfg_content['dataset']['n_classes'])
###################################### manuall modify###############################
# load model
net.to(device)
loss_function = nn.CrossEntropyLoss()
#set optimizer
###################################### manuall modify###############################
#optimizer = optim.Adam(net.parameters(),lr=cfg.lr)
optimizer = optim.SGD(net.parameters(),
lr=float(cfg_content['train']['base_lr']),
momentum=cfg_content['train']['momentum'])
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,