def __init__(self, label_file_path, weights_file_path):
#self.model = squeezenet1_1(pretrained=True)
if type(weights_file_path) != type('str'):
weights_file_path = weights_file_path.decode("utf-8")
if type(label_file_path) != type('str'):
label_file_path = label_file_path.decode("utf-8")
is_existed = os.path.exists(weights_file_path)
if not is_existed:
print("Cannot find weights file...")
return -1
print(label_file_path)
self.model = SqueezeNet(version=1.1)
self.model.load_state_dict(torch.load(weights_file_path))
self.model.eval()
self.transformation = transforms.Compose([
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
is_existed = os.path.exists(label_file_path)
if not is_existed:
print("Error: no labels file found ...")
return -1
labels = json.load(open(label_file_path))
self.class_map = {int(key): value for (key, value) in labels.items()}
print('In python: initialize successfully ...')
def squeezenet1_1(pretrained=False, **kwargs):
if pretrained:
model = SqueezeNet(version=1.1, **kwargs)
model.load_state_dict(
torch.load('./checkpoint/inception_v3_google-1a9a5a14.pth'))
return model
return SqueezeNet(version=1.1, **kwargs)
class ImgClassificationSqueezeNet:
def __init__(self, label_file_path, weights_file_path):
#self.model = squeezenet1_1(pretrained=True)
if type(weights_file_path) != type('str'):
weights_file_path = weights_file_path.decode("utf-8")
if type(label_file_path) != type('str'):
label_file_path = label_file_path.decode("utf-8")
is_existed = os.path.exists(weights_file_path)
if not is_existed:
print("Cannot find weights file...")
return -1
print(label_file_path)
self.model = SqueezeNet(version=1.1)
self.model.load_state_dict(torch.load(weights_file_path))
self.model.eval()
self.transformation = transforms.Compose([
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
is_existed = os.path.exists(label_file_path)
if not is_existed:
print("Error: no labels file found ...")
return -1
labels = json.load(open(label_file_path))
self.class_map = {int(key): value for (key, value) in labels.items()}
print('In python: initialize successfully ...')
def predict(self, image):
print("In python: predict function ...")
model = self.model
transformation = self.transformation
class_map = self.class_map
# Preprocess
image_tensor = transformation(image).float()
image_tensor = image_tensor.unsqueeze_(0)
if torch.cuda.is_available():
print("Using GPU ...")
image_tensor.cuda()
# Turn the input into a Variable
input = Variable(image_tensor)
# Predict the class of the image
output = model(input)
index = output.data.numpy().argmax()
prediction = class_map[index]
print("In python: predict done!")
return prediction
def run(self, img_data):
cv2_img = cv2.cvtColor(img_data[0], cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(cv2_img)
res = self.predict(pil_img)
return res
def get_model_by_name(model_name) -> Module:
if model_name == "vgg":
return VGG(vgg11_bn().features, num_classes=2)
if model_name == "mobilenetv2":
return MobileNetV2(num_classes=2)
if model_name == "simple":
return SimpleNetwork()
if model_name == "squeezenet":
return SqueezeNet(version=1.1, num_classes=2)
raise Exception(f"Invalid model name: {model_name}.")
if __name__ == '__main__':
device = util.get_device()
# device = torch.device('cpu')
data_loaders, data_sizes = load_data()
print(data_loaders)
print(data_sizes)
res_loss = dict()
res_top1_acc = dict()
res_top5_acc = dict()
num_classes = 100
num_epochs = 50
for name in ['3e-4 -> 1e-4', '3e-4 -> 3e-5', '3e-4 -> 0']:
model = SqueezeNet(num_classes=num_classes)
model.eval()
# print(model)
model = model.to(device)
criterion = SmoothLabelCritierion(label_smoothing=0.1)
optimizer = optim.Adam(model.parameters(), lr=3e-4, weight_decay=3e-5)
if name == '3e-4 -> 1e-4':
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs - 5, eta_min=1e-4)
elif name == '3e-4 -> 3e-5':
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs - 5, eta_min=3e-5)
else:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs - 5, eta_min=0)
warmup_scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=5, after_scheduler=lr_scheduler)
optimizer.zero_grad()
optimizer.step()
torch.save(a, 'a.pth')
# 加载为b, 存储于GPU1上(因为保存时tensor就在GPU1上)
b = torch.load('a.pth')
# 加载为c, 存储于CPU
c = torch.load('a.pth', map_location=lambda storage, loc: storage)
# 加载为d, 存储于GPU0上
d = torch.load('a.pth', map_location={'cuda:1': 'cuda:0'})
# In[40]:
torch.set_default_tensor_type('torch.FloatTensor')
from torchvision.models import SqueezeNet
model = SqueezeNet()
# module的state_dict是一个字典
model.state_dict().keys()
# Module对象的保存与加载
torch.save(model.state_dict(), 'squeezenet.pth')
model.load_state_dict(torch.load('squeezenet.pth'))
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
torch.save(optimizer.state_dict(), 'optimizer.pth')
optimizer.load_state_dict(torch.load('optimizer.pth'))
all_data = dict(optimizer=optimizer.state_dict(),
model=model.state_dict(),
info=u'模型和优化器的所有参数')
optimizer.step()
return losses
if __name__ == '__main__':
device = util.get_device()
# device = torch.device('cpu')
data_loader = load_data()
num_classes = 100
res_dict = dict()
for weight_decay in [0, 1e-3, 1e-4, 1e-5]:
# for weight_decay in [3e-5, 1e-4, 3e-4]:
model = SqueezeNet(num_classes=num_classes)
model.eval()
# print(model)
model = model.to(device)
criterion = SmoothLabelCritierion(label_smoothing=0.1)
optimizer = optim.Adam(model.parameters(),
lr=3e-4,
weight_decay=weight_decay)
losses = find_wd(data_loader, model, criterion, optimizer, device)
res_dict[str(weight_decay)] = {'loss': losses}
print('{} done'.format(weight_decay))
util.plot_loss_lr(res_dict)
print('done')
def perpare_squeezenet(squeezenet: SqueezeNet, num_classes: int):
final_conv = nn.Conv2d(512, num_classes, kernel_size=1)
squeezenet.classifier[1] = final_conv
import torch
def download(url=None):
return ''
def create():
root = download()
cifar = cleanser.ImageDataset.from_folders('cifar10', root)
def evaluate(model: cleanser.Model,
dataset: cleanser.ImageDataset,
batch_size=32):
pass
model = SqueezeNet()
transform = transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize()
])
cifar10 = cleanser.ImageDataset.objects.get(name='cifar10')
images: Iterable[cleanser.Image]
for images in cifar10.images.split('test').batches(32):
batch = torch.stack([transform(img.pillow) for img in images])
preds = model(batch)
def __init__(self, NUM_CLASSES):
super(SqueezeNet1_1, self).__init__()
self.model_name = "SqueezeNet1_1"
self.model = SqueezeNet(version=1.1, num_classes=NUM_CLASSES)