def get_mobilenet_model(n_channels, pretrained_path=None):
"""Load MobilenetV3 model with specified in and out channels"""
model = mobilenetv3_large()
if pretrained_path:
model.load_state_dict(torch.load(pretrained_path))
if n_channels == 1:
model.features[0][0].weight.data = torch.sum(
model.features[0][0].weight.data, dim=1, keepdim=True)
elif n_channels == 2:
model.features[0][0].weight.data = model.features[0][
0].weight.data[:, :2]
model.features[0][0].in_channels = n_channels
logger.debug("Mobilenet loaded. Input channels: %d", n_channels)
return model
def main():
batch_size = 64
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_iter = load_from_folder('./Train',
batch_size,
shuffle=True,
num_workers=1)
test_iter = load_from_folder('./Test',
batch_size,
shuffle=False,
num_workers=1)
lr, num_epochs = 0.001, 10
net = mobilenetv3_large(num_classes=2)
net = net.to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
print("training on ", device)
loss = torch.nn.CrossEntropyLoss()
for epoch in range(num_epochs):
train_l_sum, train_acc_sum, n, batch_count, start = 0.0, 0.0, 0, 0, time.time(
)
for X, y in train_iter:
X, y = X.to(device), y.to(device)
y_hat = net(X)
l = loss(y_hat, y)
optimizer.zero_grad()
l.backward()
optimizer.step()
train_l_sum += l.cpu().item()
train_acc_sum += (y_hat.argmax(dim=1) == y).sum().cpu().item()
n += y.shape[0]
batch_count += 1
if batch_count % 50 == 0:
print('batch %d, loss %.4f, train acc %.3f, time %.1f sec' %
(batch_count, train_l_sum / batch_count,
train_acc_sum / n, time.time() - start))
if batch_count % 1000 == 0:
test_acc = evaluate_accuracy(test_iter, net, max_batch_cnt=50)
print(
'epoch %d, loss %.4f, train acc %.3f, time %.1f sec, test acc %.3f'
% (epoch + 1, train_l_sum / batch_count, train_acc_sum / n,
time.time() - start, test_acc))
torch.save(net.state_dict(),
f'baldnet_{epoch}_{batch_count}.model')
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("image")
args = parser.parse_args()
net = mobilenetv3_large(num_classes=2)
net.load_state_dict(
torch.load("baldnet_large_0605.mdl", map_location=torch.device("cpu"))
)
net.eval()
img = cv2.imread(args.image, 1)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img_to_show = detect_and_align_face(img)
img = transforms.ToTensor()(img)
img = img.unsqueeze(0)
result = net(img)
print("result:", result)
print("秃" if result[0, 0] > result[0, 1] else "不秃")
cv2.imshow("face", cv2.cvtColor(img_to_show, cv2.COLOR_RGB2BGR))
cv2.waitKey(0)
#%%
import torch
from mobilenetv3 import mobilenetv3_large, mobilenetv3_small
#%%
net_large = mobilenetv3_large()
# net_small = mobilenetv3_small()
net_large.load_state_dict(
torch.load('pretrained/mobilenetv3-large-b4e262ea.pth'))
# net_small.load_state_dict(torch.load('pretrained/mobilenetv3-small-547c1152.pth'))
# %%
net_large.eval()
# %%
x = torch.rand(1, 3, 384, 384)
# %%
torch_res = torch.onnx._export(net_large,
x,
'pretrained/mobilenetv3-large-b4e262ea.onnx',
export_params=True,
opset_version=11)
# %%
from torchvision import models
# %%
m2 = models.mobilenet_v2(pretrained=True)
data = np.transpose(data, (2, 0, 1))
return data
top5_dict = {'0': 0, '1': 0}
top1_dict = {'0': 0, '1': 0}
ptop5_dict = {'0': 0, '1': 0}
ptop1_dict = {'0': 0, '1': 0}
if __name__ == "__main__":
src_path = "ILSVRC2012_val.txt"
if sys.argv[1] == "mobilenetv3_small":
model = mobilenetv3_small(pretrained=True)
elif sys.argv[1] == "mobilenetv3_large":
model = mobilenetv3_large(pretrained=True)
elif sys.argv[1] == "mobilenetv3_small_old":
model = mobilenetv3_small_old(pretrained=True)
elif sys.argv[1] == "mobilenetv3_large_old":
model = mobilenetv3_large_old(pretrained=True)
else:
print(
'''argv[1] must in ["mobilenetv3_large", "moiblenetv3_small", "mobilenetv3_large_old", "mobilenetv3_small_old"]'''
)
exit()
model.eval()
des_path = sys.argv[1] + ".csv"
deploy_path = sys.argv[1] + ".prototxt"
model_path = sys.argv[1] + ".caffemodel"
net = caffe.Net(deploy_path, model_path, caffe.TEST)
def build_ssd(self, backbone, size, num_classes):
mobile_layers = []
extra_layers = []
loc_layers = []
conf_layers = []
# build backbone network
if backbone == 'mobilenetv3_small':
base_model = mobilenetv3_small(num_classes=num_classes,
include_top=False)
mobile_layers += base_model.get_layers()
else:
base_model = mobilenetv3_large(num_classes=num_classes,
include_top=False)
mobile_layers += base_model.get_layers()
# build extras network on the top of the backbone
in_channels = 96
for k, v in enumerate(self.cfg['extras'][str(size)]):
extra_layers.append(
nn.Sequential(
nn.Conv2d(in_channels, v, kernel_size=1, stride=1),
nn.Conv2d(v,
v,
kernel_size=3,
stride=2,
padding=1,
groups=v),
nn.Conv2d(v, v * 2, kernel_size=1, stride=1)))
in_channels = v * 2
# build fpn and classify/regression layers
mbox = self.cfg['mbox'][str(size)]
for k, v in enumerate(self.cfg['net_source']):
loc_layers += [
nn.Conv2d(mobile_layers[v].conv._modules['0'].out_channels,
mbox[k] * 4,
kernel_size=3,
padding=1)
]
conf_layers += [
nn.Conv2d(mobile_layers[v].conv._modules['0'].out_channels,
mbox[k] * num_classes,
kernel_size=3,
padding=1)
]
for k, v in enumerate(extra_layers, 4):
loc_layers += [
nn.Conv2d(v._modules['1'].out_channels,
mbox[k] * 4,
kernel_size=3,
padding=1)
]
conf_layers += [
nn.Conv2d(v._modules['1'].out_channels,
mbox[k] * num_classes,
kernel_size=3,
padding=1)
]
return nn.ModuleList(mobile_layers), nn.ModuleList(extra_layers), \
nn.ModuleList(loc_layers), nn.ModuleList(conf_layers)
def main():
logger.info('Create data directory in which models dumped.\n')
data_dir = Path.cwd().joinpath('data')
data_dir.mkdir(exist_ok=True)
logger.info('\nInitialize MobileNetV3 and load pre-trained weights\n')
model_torch = mobilenetv3_large()
state_dict = torch.load('pretrained/mobilenetv3-large-657e7b3d.pth', map_location='cpu')
model_torch.clean_and_load_state_dict(state_dict)
logger.info('\nConvert Squeeze and Excitation modules to convert the model to a Keras model.\n')
model_torch.convert_se()
for m in model_torch.modules():
m.training = False
onnx_model_path = str(data_dir.joinpath('model.onnx'))
dummy_input = torch.randn(1, 3, 224, 224)
input_names = ['image_array']
output_names = ['category']
logger.info(f'\nExport PyTorch model in ONNX format to {onnx_model_path}.\n')
torch.onnx.export(model_torch, dummy_input, onnx_model_path,
input_names=input_names, output_names=output_names)
saved_model_dir = str(data_dir.joinpath('saved_model'))
logger.info(f'\nConvert ONNX model to Keras and save as saved_model.pb.\n')
pytorch2savedmodel(onnx_model_path, saved_model_dir)
logger.info(f'\nConvert saved_model.pb to TFLite model.\n')
tflite_model_path = str(data_dir.joinpath('model.tflite'))
tflite_model = savedmodel2tflite(saved_model_dir, tflite_model_path, quantize=False)
logger.info(f'\nConvert saved_model.pb to TFLite quantized model.\n')
tflite_quantized_model_path = str(data_dir.joinpath('model_quantized.tflite'))
tflite_quantized_model = savedmodel2tflite(saved_model_dir, tflite_quantized_model_path, quantize=True)
logger.info("\nCompare PyTorch model's outputs and TFLite models' outputs.\n")
num_same_outputs = 0
image_path_list = list(Path('tools').glob('*.jpg'))
for path in image_path_list:
input_array = load_and_preprocess_image(str(path))
input_tensor = torch.from_numpy(input_array)
torch_output = model_torch(input_tensor).data.numpy().reshape(-1, )
tflite_output = get_tflite_outputs(input_array.transpose((0, 2, 3, 1)), tflite_model).reshape(-1, )
logger.info(f'PyTorch - first 5 items: {torch_output[:5]}')
logger.info(f'TFLite - first 5 items: {tflite_output[:5]}')
torch_output_index = np.argmax(torch_output)
tflite_output_index = np.argmax(tflite_output)
logger.info(f'PyTorch - argmax index: {torch_output_index}')
logger.info(f'TFLite - argmax index: {tflite_output_index}\n')
if torch_output_index == tflite_output_index:
num_same_outputs += 1
logger.info(f'# of matched outputs: {num_same_outputs} / {len(image_path_list)}\n')
logger.info("\nCompare PyTorch model's outputs and TFLite quantized models' outputs.\n")
num_same_outputs = 0
image_path_list = list(Path('tools').glob('*.jpg'))
for path in image_path_list:
input_array = load_and_preprocess_image(str(path))
input_tensor = torch.from_numpy(input_array)
torch_output = model_torch(input_tensor).data.numpy().reshape(-1, )
tflite_output = get_tflite_outputs(input_array.transpose((0, 2, 3, 1)), tflite_quantized_model).reshape(-1, )
logger.info(f'PyTorch - first 5 items: {torch_output[:5]}')
logger.info(f'TFLite - first 5 items: {tflite_output[:5]}')
torch_output_index = np.argmax(torch_output)
tflite_output_index = np.argmax(tflite_output)
logger.info(f'PyTorch - argmax index: {torch_output_index}')
logger.info(f'TFLite - argmax index: {tflite_output_index}\n')
if torch_output_index == tflite_output_index:
num_same_outputs += 1
logger.info(f'# of matched outputs: {num_same_outputs} / {len(image_path_list)}\n')