def vgg11_bn(num_classes=1000, pretrained='imagenet'):
"""VGG 11-layer model (configuration "A") with batch normalization
"""
model = models.vgg11_bn(pretrained=False)
if pretrained is not None:
settings = pretrained_settings['vgg11_bn'][pretrained]
model = load_pretrained(model, num_classes, settings)
return model
vgg11_layers = get_vgg_layers(vgg11_config, batch_norm=True) print(vgg11_layers) # %% OUTPUT_DIM = 7 model = toy_VGG(vgg11_layers, OUTPUT_DIM) print(model) # IN_FEATURES = pretrained_model.classifier[-1].in_features # final_fc = nn.Linear(IN_FEATURES, OUTPUT_DIM) # pretrained_model.classifier[-1] = final_fc # print(pretrained_model.classifier) # %% import torchvision.models as models pretrained_model = models.vgg11_bn(pretrained=True) # %% pretrained_model.classifier[-1] # %% IN_FEATURES = pretrained_model.classifier[-1].in_features final_fc = nn.Linear(IN_FEATURES, OUTPUT_DIM) # %% pretrained_model.classifier[-1] = final_fc # %% print(pretrained_model.classifier) # %%
def test_vgg11_bn(self):
process_model(models.vgg11_bn(self.pretrained), self.image, _C_tests.forward_vgg11bn, 'VGG11BN')
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet201(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs,num_classes)
input_size = 299
elif model_name == 'se_resnext50_32x4d':
"se_resnext 50 32x4d"
if use_pretrained==True:
model_ft = pretrainedmodels.__dict__[model_name](pretrained='imagenet')
dim_feats = model_ft.last_linear.in_features # =2048
model_ft.last_linear = nn.Linear(dim_feats, num_classes)
else:
#to do
model_ft = pretrainedmodels.__dict__[model_name](pretrained=None)
dim_feats = model_ft.last_linear.in_features # =2048
model_ft.last_linear = nn.Linear(dim_feats, num_classes)
input_size = 224
elif model_name == 'se_resnet50':
if use_pretrained==True:
model_ft = pretrainedmodels.__dict__[model_name](pretrained='imagenet')
dim_feats = model_ft.last_linear.in_features # =2048
model_ft.last_linear = nn.Linear(dim_feats, num_classes)
else:
#to do
model_ft = pretrainedmodels.__dict__[model_name](pretrained=None)
dim_feats = model_ft.last_linear.in_features # =2048
model_ft.last_linear = nn.Linear(dim_feats, num_classes)
input_size=224
else:
if use_pretrained==True:
model_ft = pretrainedmodels.__dict__[model_name](pretrained='imagenet')
dim_feats = model_ft.last_linear.in_features # =2048
model_ft.last_linear = nn.Linear(dim_feats, num_classes)
else:
#to do
model_ft = pretrainedmodels.__dict__[model_name](pretrained=None)
dim_feats = model_ft.last_linear.in_features # =2048
model_ft.last_linear = nn.Linear(dim_feats, num_classes)
input_size=224
return model_ft, input_size
def taehong(image_name):
# Creat Net
net_face = models.vgg11_bn(pretrained=True)
net_first = models.vgg11_bn(pretrained=True)
old_classifier1 = list(
net_face.classifier.children()) # get the classifier part alone
old_classifier1.pop() # remove the last layer
old_classifier1.append(nn.Linear(4096, 7)) # add a new layer
net_face.classifier = nn.Sequential(
*old_classifier1) # attach it to the original vgg mode
old_classifier2 = list(
net_first.classifier.children()) # get the classifier part alone
old_classifier2.pop() # remove the last layer
old_classifier2.append(nn.Linear(4096, 6)) # add a new layer
net_first.classifier = nn.Sequential(
*old_classifier2) # attach it to the original vgg mode
# Load Net
net_face.load_state_dict(torch.load('faceClssifcation_99.pth'))
net_first.load_state_dict(torch.load('firstFeelingClssifcation_88.pth'))
# Creast class
face_classes = ('조류상', '고양이상', '공룡상', '강아지상', '여우상', '개구리상', '말상')
first_class = ('감정이 과한', '과묵한', '불만 많은', '사연있는', '호감형', '화끈한')
# Data preprocessing
resize_data = data_preprocessing(image_name)
# Read image
image = image_loader(resize_data)
# Output
output_face = net_face(image)
output_first = net_first(image)
_, face_preds = torch.max(output_face.data, 1)
_, first_preds = torch.max(output_first.data, 1)
face = face_classes[np.asarray(face_preds)[0]]
first = first_class[np.asarray(first_preds)[0]]
# microsoft lens
face_api_url = 'https://westcentralus.api.cognitive.microsoft.com/face/v1.0/detect'
headers = {
'ocp-apim-subscription-key': '71199991076143889c95226508565943',
'Content-Type': "application/octet-stream",
'cache-control': "no-cache",
}
#header는 개인 아이디 신청하여 바꾸시면 됩니다
params = {
'returnFaceAttributes': 'gender,age,emotion',
}
data = open(image_name, 'rb').read()
response = requests.post(face_api_url,
params=params,
headers=headers,
data=data)
faces = response.json()
dic1 = dict(faces[0])
dic2 = dic1['faceAttributes']
age = dic2['age']
gender = dic2['gender']
dic3 = dict(dic2['emotion'])
inverse = [(value, key) for key, value in dic3.items()]
emotion = max(inverse)
output_emotion = emotion[1]
title = creat_title(age, output_emotion, gender, face)
return age, output_emotion, gender, face, first, title
def trainDec():
testSplit = split
nSplit = 5
listSplit = []
for i in range(nSplit):
if i!=testSplit :
listSplit.append(i)
print(listSplit)
model_ft = None
batch_size = 8
save_name_ori = ""
image_size = 224
multi_gpu = False
resume = True
num_classes = 2
num_epochs = 1000
freezeInternal = False
#Intialize the model for this run
model_sel = type #0 : senet, 1 : resnet (other) , 3 : senet50, 4 : resnet50(offi), 5. short .
use_internal_n = True
if model_sel == 0 :
batch_size =20
import senet as SENet
model_ft = SENet.senet50(num_classes=8631)
save_name_ori = 'senet50_other.pkl'
elif model_sel == 1 :
batch_size = 6 #60
import resnet as ResNet
#model_ft = senet154(num_classes=1000, pretrained='imagenet')
model_ft = ResNet.resnet50(num_classes=8631)
save_name_ori = 'resnet50_other.pkl'
elif model_sel == 2 :
batch_size =32
'''import senet50_256_pytorch as SENetOri
save_name_ori = 'senet50_256_pytorch'
#model_ft = imp.load_source('MainModel', save_name_ori+'.py')
model_ft = SENetOri.KitModel(save_name_ori+'.pth')'''
import senet50_256_fc as SENetOri
save_name_ori = 'senet50_256_pytorch'
#model_ft = imp.load_source('MainModel', save_name_ori+'.py')
model_ft = SENetOri.KitModel(save_name_ori+".pth")
num_inter = 256
elif model_sel == 3 :
batch_size =32
'''import senet50_256_pytorch as SENetOri
save_name_ori = 'senet50_256_pytorch'
#model_ft = imp.load_source('MainModel', save_name_ori+'.py')
model_ft = SENetOri.KitModel(save_name_ori+'.pth')'''
import resnet50_128_fc as ResNetOri
save_name_ori = 'resnet50_128_pytorch'
#model_ft = imp.load_source('MainModel', save_name_ori+'.py')
model_ft = ResNetOri.KitModel(save_name_ori+".pth")
num_inter = 128
elif model_sel == 4 :
batch_size =8
import resnet50_128_fc as ResNetOri
save_name_ori = 'resnet50_128_pytorch'
#model_ft = imp.load_source('MainModel', save_name_ori+'.py')
model_ft = ResNetOri.KitModel(save_name_ori+".pth")
num_inter = 128
elif model_sel == 5 :
batch_size =20
model_ft = models.vgg11_bn(pretrained=True)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
save_name_ori = 'vgg11.pt'
elif model_sel == 6 :
batch_size = 50
model_ft = models.resnet152(pretrained = True) #50 or 152
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
save_name_ori = 'resnet152.pt'
if runServer :
batch_size = int(batch_size + round(truediv(batch_size, 4)))
# Just normalization for validation
transform =transforms.Compose([
transforms.Resize((image_size,image_size)),
#transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
#transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
print(model_ft)
print('batch_size : ',batch_size)
save_name = 'AFEW-VA-'+str(testSplit)+"-"+save_name_ori
err_file = curDir+save_name+"AFEWPART2.txt"
set_parameter_requires_grad(model_ft, freezeInternal)
if multi_gpu :
batch_size*=2 #'300W-Train','Menpo_Challenge/2D','300W_LP'
ID = AFEWVA(["AFEW-VA-PP"], None, True, image_size, transform, True, False, 1,split=True, nSplit = nSplit,listSplit=listSplit,wHeatmap=False,isVideo=False,seqLength=5)
#ID =FacialLandmarkDataset(['300W-Train','Menpo_Challenge/2D','300W_LP'], None, True, image_size, transform, use_internal_n, True, 1)
dataloader = torch.utils.data.DataLoader(dataset = ID, batch_size = batch_size, shuffle = True)
VD = AFEWVA(["AFEW-VA-PP"], None, True, image_size, transform, True, False, 1,split=True, nSplit = nSplit,listSplit=[testSplit],wHeatmap=False,isVideo=False,seqLength=5)
dataloaderV = torch.utils.data.DataLoader(dataset = VD, batch_size = batch_size, shuffle = True)
if resume :
pretrained_dict = torch.load(curDir+'t-models/'+save_name)
model_dict = model_ft.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model_ft.load_state_dict(pretrained_dict)
model_ft = model_ft.cuda()#to(device)
params_to_update = model_ft.parameters()
print("Params to learn ")
if freezeInternal :
params_to_update = []
for name,param in model_ft.named_parameters():
if param.requires_grad == True :
params_to_update.append(param)
print("\t",name)
else :
for name,param in model_ft.named_parameters() :
if param.requires_grad == True :
print("\t",name)
optimizer_ft = optim.SGD(params_to_update,lr=0.001, momentum = .9)
#loss
criterion = nn.MSELoss()
#####################Training
x1,lbl,ldmrk = next(iter(dataloader))
x1 = x1.cuda();lbl = lbl.cuda();
toDisplay = 4
f = open(err_file,'w+')
f.write("err : ")
f.close()
val_acc_history = []
#best_model_wts = copy.deepcopy(model.state_dict())
lowest_loss = 99999
for epoch in range(num_epochs) :
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-'*10)
f = open(err_file,'a')
f.write('Epoch {}/{}'.format(epoch, num_epochs - 1))
f.close()
running_loss = 0
running_corrects = 0
#for rinputs, rlabels in dataloaders :
for x,(rinputs, rlabels,rldmrk) in enumerate(dataloader,0) :
model_ft.train()
inputs = rinputs.cuda()#to(device)
labels = rlabels.cuda()#to(device)
#zero the parameter gradients
optimizer_ft.zero_grad()
with torch.set_grad_enabled(True) :
outputs = model_ft(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer_ft.step()
#statistics
running_loss += loss.item() * inputs.size(0)
print("{}/{} loss : {}".format(x,int(len(dataloader.dataset)/batch_size),loss.item()))
f = open(err_file,'a')
f.write("{}/{} loss : {}\n".format(x,int(len(dataloader.dataset)/batch_size),loss.item()))
f.close()
epoch_loss = running_loss / len(dataloader.dataset)
print('Loss : {:.4f}'.format(epoch_loss))
f = open(err_file,'a')
f.write('Loss : {:.4f}'.format(epoch_loss))
f.close()
#Deep copy the model_ft
if epoch%2 == 0 :#epoch_loss < lowest_loss :
lowest_loss = lowest_loss
print("outoput : ",outputs[0])
print("labels : ",labels[0])
if multi_gpu :
torch.save(model_ft.module.state_dict(),curDir+'t-models/'+save_name)
else :
torch.save(model_ft.state_dict(),curDir+'t-models/'+save_name)
if True :
listValO = []
listAroO = []
listValL = []
listAroL = []
tvo = [];tao=[];tvl = []; tal = [];
anyDiffer = False
for x,(rinputs, rlabels,rldmrk) in enumerate(dataloaderV,0) :
model_ft.eval()
inputs = rinputs.cuda()#to(device)
labels = rlabels.cuda()#to(device)
with torch.set_grad_enabled(False) :
outputs = model_ft(inputs)
#print(x,',',int(truediv(len(VD),batch_size)),outputs[:2], labels[:2],outputs[:,0].shape[0],outputs.shape)
if outputs[:,0].shape[0] != batch_size : #in case the batch size is differ, usually at end of iter
anyDiffer = True
print('differ')
tvo.append(outputs[:,0].detach().cpu())
tao.append(outputs[:,1].detach().cpu())
tvl.append(labels[:,0].detach().cpu())
tal.append(labels[:,1].detach().cpu())
else :
print('equal')
listValO.append(outputs[:,0].detach().cpu())
listAroO.append(outputs[:,1].detach().cpu())
listValL.append(labels[:,0].detach().cpu())
listAroL.append(labels[:,1].detach().cpu())
est_V = np.asarray(torch.stack(listValO)).flatten()
est_A = np.asarray(torch.stack(listAroO)).flatten()
gt_V = np.asarray(torch.stack(listValL)).flatten()
gt_A = np.asarray(torch.stack(listAroL)).flatten()
if anyDiffer :
est_Vt = np.asarray(torch.stack(tvo)).flatten()
est_At = np.asarray(torch.stack(tao)).flatten()
gt_Vt = np.asarray(torch.stack(tvl)).flatten()
gt_At = np.asarray(torch.stack(tal)).flatten()
#now concatenate
est_V = np.concatenate((est_V,est_Vt))
est_A = np.concatenate((est_A,est_At))
gt_V = np.concatenate((gt_V,gt_Vt))
gt_A = np.concatenate((gt_A,gt_At))
print(est_V.shape, gt_V.shape)
mseV = calcMSE(est_V, gt_V)
mseA = calcMSE(est_A, gt_A)
corV = calcCOR(est_V, gt_V)
corA = calcCOR(est_A, gt_A)
iccV = calcICC(est_V, gt_V)
iccA = calcICC(est_A, gt_A)
iccV2 = calcICC(gt_V, gt_V)
iccA2 = calcICC(gt_A, gt_A)
print('epoch : ',epoch, ', epoch_loss : ',epoch_loss)
print('MSEV : ',mseV, ', CORV : ',corV,', ICCV : ',iccV,', ICCV2 : ',iccV2)
print('MSEA : ',mseA, ', CORA : ',corA,', ICCA : ',iccA,', ICCA2 : ',iccA2)
with open('AFEW-VA-RESNET-'+str(testSplit)+'-resPART2.csv',mode='a') as fcsv:
fcsv.write(str(epoch)+','+str(epoch_loss)+','+str(mseV)+','+str(corV)+','+str(iccV)+','+str(iccV2)+','+str(mseA)+','+str(corA)+','+str(iccA)+','+str(iccA2)+'\n')
print('Best val Acc: {:4f}'.format(lowest_loss))
root=conf["dataset"]["validate"], transform=transform_list)
validate_loader = torch.utils.data.DataLoader(
validate_set,
batch_size=conf["parameters"]["batch_size"],
shuffle=True,
num_workers=4)
assert train_set.class_to_idx == validate_set.class_to_idx
# Model
model_name = args.model if args.model else conf["parameters"]["model_name"]
print(f">> Building {model_name} model...")
model_dict = {
"shufflenet": shufflenetv2.shufflenet_v2_x1_0(pretrained=True),
"mobilenet": mobilenet.mobilenet_v2(pretrained=True),
"vgg": vgg11_bn(pretrained=True),
"resnet": resnet18(pretrained=True),
"crnn": crnn.CRNN()
}
model = model_dict[model_name]
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=conf["parameters"]["learning_rate"],
momentum=0.9,
weight_decay=5e-4)
def train():
best_train_acc = 0
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
model_ft = None
input_size = 0
if model_name == "resnet":
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_fltrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fltrs, num_classes)
input_size = 224
elif model_name == "alexnet":
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_fltrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_fltrs, num_classes)
input_size = 224
elif model_name == "vgg":
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_fltrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_fltrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_fltrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_fltrs, num_classes)
input_size = 224
elif model_name == "inception":
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_fltrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_fltrs, num_classes)
# Handle the primary net
num_fltrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fltrs, num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
# Ignore ssl certification (prevent error for some users)
ssl._create_default_https_context = ssl._create_unverified_context
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "resnet152":
model_ft = models.resnet152(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
elif model_name == "xception":
""" Xception
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = xception.xception(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
elif model_name == "fleming_v1":
""" Fleming Model
Custom model created by team Fleming
"""
model_ft = fleming.FlemingModel_v1(num_classes=196)
input_size = 224
elif model_name == "fleming_v2":
""" Fleming Model
Custom model created by team Fleming
"""
model_ft = fleming.FlemingModel_v2(num_classes=196)
input_size = 224
elif model_name == "fleming_v3":
""" Fleming Model
Custom model created by team Fleming
"""
model_ft = fleming.FlemingModel_v3(num_classes=196)
input_size = 224
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def initialize_model(model_name, use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 224
if model_name == "resnet18":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
elif model_name == "resnet34":
""" Resnet34
"""
model_ft = models.resnet34(pretrained=use_pretrained)
elif model_name == "resnet50":
""" Resnet50
"""
model_ft = models.resnet50(pretrained=use_pretrained)
elif model_name == "resnet101":
""" Resnet101
"""
model_ft = models.resnet101(pretrained=use_pretrained)
elif model_name == "resnet152":
""" Resnet152
"""
model_ft = models.resnet152(pretrained=use_pretrained)
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
elif model_name == "vgg11_bn":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
elif model_name == "squeezenet1_0":
""" Squeezenet1_0
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
elif model_name == "densenet121":
""" Densenet121
"""
model_ft = models.densenet121(pretrained=use_pretrained)
elif model_name == "densenet169":
""" Densenet169
"""
model_ft = models.densenet169(pretrained=use_pretrained)
elif model_name == "densenet201":
""" Densenet201
"""
model_ft = models.densenet201(pretrained=use_pretrained)
elif model_name == "densenet161":
""" Densenet161
"""
model_ft = models.densenet161(pretrained=use_pretrained)
elif model_name == "inception_v3":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
data = RetDataset(csv_file, image_root, transform=transform)
if alexnet:
model = models.alexnet(pretrained=True)
set_parameter_requires_grad(model, True)
num_ftrs = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_ftrs, 5)
feature_extractor = models.alexnet(pretrained=True)
# feature_extractor.load_state_dict(torch.load(load_weights_dir, map_location='cpu')) if not frozen
beforelabel = nn.Sequential(*list(model.classifier.children())[:-1])
feature_extractor.classifier = beforelabel
elif vgg:
model = models.vgg11_bn(pretrained=True)
set_parameter_requires_grad(model, True)
num_ftrs = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_ftrs, 5)
feature_extractor = models.vgg11_bn(pretrained=True)
beforelabel = nn.Sequential(*list(model.classifier.children())[:-2])
feature_extractor.classifier = beforelabel
elif resnet:
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 5)
feature_extractor = nn.Sequential(*list(model.children())[:-1])
print(feature_extractor)
def vgg11_bn(num_classes, pretrained=False):
return models.vgg11_bn(pretrained=pretrained, num_classes=num_classes)
def __init__(self):
super(Vgg11bn, self).__init__()
model_vgg11bn = models.vgg11_bn(pretrained=True)
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=False):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
# if model_name == "resnet":
# """ Resnet34
# """
# model_ft = tm.resnet34(pretrained=use_pretrained)
# set_parameter_requires_grad(model_ft, feature_extract)
# num_ftrs = model_ft.fc.in_features
# model_ft.fc = nn.Linear(num_ftrs, num_classes)
# input_size = 224
elif model_name == "regnet":
""" regnet
regnety_040， regnety_080， regnety_160
"""
model_ft = tm.regnety_040(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.get_classifier().in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "efficientnet_b2":
"""
efficientnet_b2 256, efficientnet_b3 288, efficientnet_b4 320
"""
model_ft = tm.efficientnet_b2(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.get_classifier().in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 256
elif model_name == "efficientnet_b3":
"""
efficientnet_b2 256, efficientnet_b3 288, efficientnet_b4 320
"""
model_ft = tm.efficientnet_b3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.get_classifier().in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 288
elif model_name == "efficientnet_b4":
"""
efficientnet_b2 256, efficientnet_b3 288, efficientnet_b4 320
"""
model_ft = tm.efficientnet_b4(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.get_classifier().in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 320
elif model_name == "vit":
""" vit
"""
model_ft = tm.vit_tiny_patch16_224(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.head.in_features
model_ft.head = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "pit":
""" pit
pit_xs_224, pit_s_224
"""
model_ft = tm.pit_xs_224(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.get_classifier().in_features
model_ft.head = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "deit":
""" deit
deit_small_patch16_224, deit_base_patch16_224
"""
model_ft = tm.deit_small_patch16_224(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.head.in_features
model_ft.head = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "mixer":
""" mixer
"""
model_ft = tm.mixer_b16_224(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.head.in_features
model_ft.head = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "swin-vit":
""" swin-vit
tm.swin_tiny_patch4_window7_224, tm.swin_small_patch4_window7_224
"""
model_ft = tm.swin_small_patch4_window7_224(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.head.in_features
model_ft.head = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs,num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def vgg11(pre):
return children(vgg11_bn(pre))[0]
layers.append(layer)
in_channels = c
layers = nn.Sequential(*layers)
return layers
output_dim = 10
vgg11_config = [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M']
features = get_vgg_layers(vgg11_config, batch_norm=True)
model = VGG(features, output_dim)
print(model)
import torchvision.models as vision_models
vgg_pretrained = vision_models.vgg11_bn(pretrained=True)
in_features = vgg_pretrained.classifier[-1].in_features
vgg_pretrained.classifier[-1] = nn.Linear(in_features, output_dim)
""""""
print(vgg_pretrained)
vgg_pretrained_params = list(vgg_pretrained.named_parameters())
model_params = list(model.named_parameters())
assert len(vgg_pretrained_params) == len(model_params)
for i in range(len(vgg_pretrained_params)):
assert vgg_pretrained_params[i][0] == model_params[i][0]
return pretrained
def forward(self, x):
"""
Forward propages the network given an input batch
:param x: Inputs x (b, c, h, w)
:return: preds (b, num_classes)
"""
return self.model(x)
# def reset_parameters(self):
# """
# Re-initialize the network parameters.
# """
# for item in self.layer_dict.children():
# try:
# item.reset_parameters()
# except:
# pass
# self.logit_linear_layer.reset_parameters()
# conv_net = ConvolutionalNetwork(10)
pretrained = models.vgg11_bn(pretrained=False)
for i, j in pretrained.named_parameters():
print(i, j)
# pretrained.fc = nn.Linear(2048,self.num_output_classes)
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs,num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def _get_model_and_layer(self, model_name, layer):
""" Internal method for getting layer from model
:param model_name: model name such as 'resnet-18'
:param layer: layer as a string for resnet-18 or int for alexnet
:returns: pytorch model, selected layer
"""
if model_name.startswith(
'resnet') and not model_name.startswith('resnet-'):
model = getattr(models, model_name)(pretrained=True)
if layer == 'default':
layer = model._modules.get('avgpool')
self.layer_output_size = self.RESNET_OUTPUT_SIZES[model_name]
else:
layer = model._modules.get(layer)
return model, layer
elif model_name == 'resnet-18':
model = models.resnet18(pretrained=True)
if layer == 'default':
layer = model._modules.get('avgpool')
self.layer_output_size = 512
else:
layer = model._modules.get(layer)
return model, layer
elif model_name == 'alexnet':
model = models.alexnet(pretrained=True)
if layer == 'default':
layer = model.classifier[-2]
self.layer_output_size = 4096
else:
layer = model.classifier[-layer]
return model, layer
elif model_name == 'vgg':
# VGG-11
model = models.vgg11_bn(pretrained=True)
if layer == 'default':
layer = model.classifier[-2]
self.layer_output_size = model.classifier[
-1].in_features # should be 4096
else:
layer = model.classifier[-layer]
return model, layer
elif model_name == 'densenet':
# Densenet-121
model = models.densenet121(pretrained=True)
if layer == 'default':
layer = model.features[-1]
self.layer_output_size = model.classifier.in_features # should be 1024
else:
raise KeyError('Un support %s for layer parameters' %
model_name)
return model, layer
else:
raise KeyError('Model %s was not found' % model_name)
def initialize_model(model_name,
num_classes,
feature_extract=False,
use_pretrained=True):
model_ft = None
input_size = 0
if model_name == 'resnet18':
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'resnet101':
""" Resnet101
"""
model_ft = models.resnet101(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'resnext':
""" ResNext50_32x4d
"""
model_ft = models.resnext50_32x4d(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'alexnet':
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'vgg':
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'squeezenet':
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == 'inception':
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
else:
logger.warning("Invalid model name, exiting...")
exit()
return model_ft, input_size
def initialize_model(model_name, num_classes, feature_extract, use_pretrained = True):
model_ft = None
input_size = 0
if model_name == 'resnet' :
model_ft = models.resnet152(pretrained = use_pretrained) #50 or 152
'''
in case, feature_extract is true, set_parameter_requires_grad will set all grad parameters to false
and because after this operation a new layer (the output one) is added, these new parameters will have grad as true
'''
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'alexnet' :
model_ft = models.alexnet(pretrained = use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'vgg':
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == 'squeezenet' :
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,num_classes, kernel_size=(1,1), stride = (1,1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == 'densenet' :
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == 'inception' :
#inception v3
model_ft = models.inception_v3(pretrained = use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
#Aux net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
#Primary Net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
else :
print('invalid model name')
exit()
return model_ft, input_size
def train(gpu,args):
img_transforms = Compose([
CenterCrop((224, 224)),
ToTensor(),
Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
classes = ['BlowDryHair', 'Typing', 'BandMarching', 'BoxingSpeedBag', 'PlayingCello', 'PlayingDaf', 'PlayingDhol','PlayingFlute', 'PlayingSitar', 'BrushingTeeth']
# classes = ['LongJump',
# 'CricketBowling',
# 'CuttingInKitchen',
# 'ApplyEyeMakeup',
# 'FrontCrawl',
# 'HammerThrow',
# 'Bowling',
# 'BodyWeightSquats',
# 'SkyDiving',
# 'StillRings',
# 'PlayingCello',
# 'CricketShot',
# 'Archery',
# 'CliffDiving',
# 'Knitting',
# 'BoxingPunchingBag',
# 'SoccerPenalty',
# 'Hammering',
# 'SumoWrestling',
# 'FrisbeeCatch',
# 'ShavingBeard',
# 'PlayingDaf',
# 'FloorGymnastics',
# 'BasketballDunk',
# 'WritingOnBoard',
# 'PlayingFlute',
# 'BabyCrawling',
# 'PlayingDhol',
# 'BalanceBeam',
# 'BrushingTeeth',
# 'ApplyLipstick',
# 'ParallelBars',
# 'PlayingSitar',
# 'HandstandWalking',
# 'FieldHockeyPenalty',
# 'HeadMassage',
# 'MoppingFloor',
# 'Shotput',
# 'BandMarching',
# 'Haircut',
# 'TableTennisShot',
# 'Typing',
# 'BoxingSpeedBag',
# 'UnevenBars',
# 'Surfing',
# 'HandstandPushups',
# 'Rafting',
# 'IceDancing',
# 'BlowDryHair',
# 'WallPushups',
# 'BlowingCandles']
classes_num = len(classes)
torch.cuda.set_device(gpu)
#### Set the dataset
rank = args.nr * args.gpus + gpu
print('rank',rank)
seed = rank+args.seed
print('seed',seed)
init_seeds(seed)
train_dataset = UCF101Frame_Sound_Dataset(classes, args.train_path,img_transform = img_transforms,aud_transform = None)
val_dataset = UCF101Frame_Sound_Dataset_test(classes, args.test_path,img_transform = img_transforms,aud_transform = None)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=args.world_size,
rank=rank,
shuffle=True
)
train_labels = train_dataset.all_labels
print('true labels',Counter(train_labels))
##generate noise label
if args.y_noise_generate == True:
result_path = 'results/'+args.model + '/noise_type_{}_noise_percent{}_seed{}_epochs{}'.format(args.noise_type,args.noise_percent,args.seed,args.epoches)
if args.noise_type == 'symmetry':
corruption_matrix = uniform_mix_C(args.noise_percent,classes_num)
elif args.noise_type == 'asymmetry_balance':
corruption_matrix = flip_labels_C_balanced(args.noise_percent, classes_num,args.seed)
elif args.noise_type == 'asymmetry':
corruption_matrix = flip_labels_C(args.noise_percent, classes_num,args.seed)
elif args.noise_type == 'asymmetry2':
corruption_matrix = flip_labels_C_two(args.noise_percent, classes_num,args.seed)
print('corruption_matrix',corruption_matrix)
#np.random.seed(args.seed)
for i in range(len(train_labels)):
train_dataset.all_labels[i] = np.random.choice(classes_num, p=corruption_matrix[train_labels[i]])
print('noise labels', Counter(train_dataset.all_labels))
else:
print('No noise labels')
result_path = 'results/'+args.model + '/NoNoise_seed{}_epochs{}'.format(args.seed,args.epoches)
train_dataloader = DataLoaderX(train_dataset, batch_size=args.batch_size, shuffle=False, num_workers = 4,pin_memory=True, sampler=train_sampler)
val_dataloader = DataLoaderX(val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True,sampler=None)
###set the distribution
dist.init_process_group(
backend='nccl',
init_method='env://',
world_size=args.world_size,
rank=rank
)
if args.model == 'Vision-Subnet':
model = ImageSubNet(classes_num).cuda()
elif args.model == 'VGG16':
model = models.vgg11_bn(pretrained=False).cuda()
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) #实现多卡BN
model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu], output_device=gpu)
crossEntropy = nn.CrossEntropyLoss().cuda()
#optim = Adam(model.parameters(), lr = 1e-5 , weight_decay = 1e-5,amsgrad=True)
optim = torch.optim.SGD(model.parameters(), lr = 5e-5, momentum=0.9)
#每16个周期降低6%
scheduler = lr_scheduler.StepLR(optim,16,0.94)
scaler = torch.cuda.amp.GradScaler()
total_batch = len(train_dataloader.dataset) / args.batch_size / args.world_size
train_loss_sequence = []
train_acc_sequence = []
val_loss_sequence = []
val_acc_sequence = []
max_val_acc = 0
epochs_no_improve = 0
frame_level_pred, video_level_label = get_vedio_dict(classes,args.splite_path)
if os.path.exists(result_path) is not True:
os.makedirs(result_path+'/model/')
os.makedirs(result_path + '/plot/')
os.makedirs(result_path + '/log/')
for epoch in range(args.epoches):
model.train()
batch_loss_sequence = []
batch_acc_sequence = []
for batch_idx, (img,_,label) in enumerate(train_dataloader):
optim.zero_grad()
img = img.float()
img = img.cuda()
label = label.cuda()
with torch.cuda.amp.autocast():
out = model(img)
loss = crossEntropy(out, label)
scaler.scale(loss).backward()
scaler.step(optim)
scaler.update()
_, predict = out.max(1)
correct = (predict.data == label.data).sum() * 1.0
batch_acc = correct / len(label)
batch_loss_sequence.append(loss.cpu().detach().numpy())
batch_acc_sequence.append(batch_acc.cpu().detach().numpy())
if (batch_idx + 1) % 5 == 0:
print('# Epoch %3d: batch_idx/ total_batch %3d/%3d: train_loss: %.6f batch_acc: %.6f' % (epoch + 1, batch_idx, total_batch, loss,batch_acc))
scheduler.step() # adjust lr
if (epoch + 1) % 1 == 0:
if rank == 0:
val_loss, val_acc = test(model,val_dataloader,frame_level_pred, video_level_label,args)
train_loss_sequence.append(np.mean(batch_loss_sequence))
train_acc_sequence .append(np.mean(batch_acc_sequence))
val_loss_sequence.append(float(val_loss))
val_acc_sequence.append(float(val_acc))
if val_acc > max_val_acc:
epochs_no_improve = 0
max_val_acc = val_acc
#torch.save(model.state_dict(),result_path+'/model/best.model')
else:
epochs_no_improve += 1
if epochs_no_improve == args.n_epochs_stop:
print('Early Stopping!')
break
print('# Epoch %3d: train_loss: %.6f val_loss: %.6f val_acc: %.6f best: %.6f' % (epoch + 1, loss, val_loss, val_acc, max_val_acc))
if rank == 0:
print('train_loss_sequence',train_loss_sequence)
print('train_acc_sequence', train_acc_sequence)
print('val_loss_sequence', val_loss_sequence)
print('val_acc_sequence', val_acc_sequence)
plot_process(result_path + '/plot/train_val.jpg',list(range(len(train_loss_sequence))),train_loss_sequence,train_acc_sequence,val_loss_sequence,val_acc_sequence)
if (modelName == 'alexnet'):
model = models.alexnet(pretrained=True, progress=False)
elif (modelName == 'resnet18'):
model = models.resnet18(pretrained=True, progress=False)
elif (modelName == 'resnet34'):
model = models.resnet34(pretrained=True, progress=False)
elif (modelName == 'resnet50'):
model = models.resnet50(pretrained=True, progress=False)
elif (modelName == 'resnet101'):
model = models.resnet101(pretrained=True, progress=False)
elif (modelName == 'resnet152'):
model = models.resnet152(pretrained=True, progress=False)
elif (modelName == 'vgg11'):
model = models.vgg11(pretrained=True, progress=False)
elif (modelName == 'vgg11_bn'):
model = models.vgg11_bn(pretrained=True, progress=False)
elif (modelName == 'squeezenet1_0'):
model = models.squeezenet1_0(pretrained=True, progress=False)
elif (modelName == 'squeezenet1_1'):
model = models.squeezenet1_1(pretrained=True, progress=False)
elif (modelName == 'densenet161'):
model = models.densenet161(pretrained=True, progress=False)
elif (modelName == 'shufflenet_v2_x0_5'):
model = models.shufflenet_v2_x0_5(pretrained=True, progress=False)
elif (modelName == 'mobilenet_v2'):
model = models.mobilenet_v2(pretrained=True, progress=False)
elif (modelName == 'mnasnet1_0'):
model = models.mnasnet1_0(pretrained=True, progress=False)
elif (modelName == 'googlenet'):
model = models.googlenet(pretrained=True, progress=False)
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
"""
Source: https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
Initialize these variables which will be set in this if statement. Each of these
variables is model specific.
:param str model_name: model to be loaded
:param int num_classes: number of classes
:param bool feature_extract: deactivate gradients
:param bool use_pretrained: load pretrained weights
:return: pretrained model, input_size
"""
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3, Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def _make_model(base_model=_std_base_model,
hidden_1_out=_std_hidden_1_out,
hidden_2_out=_std_hidden_2_out,
dropout_rate=_std_dropout_rate,
device=_std_device):
'''
Description:
============
Helper function.
Creates a customizable model based on selected torchvision models
by replacing the classifier with a linear two hidden layer relu neutral network
for cross-entropy with (given the inputs from the feature extractor of the base model
and given the necessary outputs to match number of the classes)
customizable hidden units in layer 1 and layer 2 as well as
customizable dropout rate
Keyword Arguments:
- base_model:
Name of the pytorch model as string (must be in available base models)
(default: 'alexnet')
- hidden_1_out/hidden_2_out:
Number of hidden outputs in layer 1/layer 2 as int (default: 256)
- dropout_rate:
dropout in layer 1 and layer 2 as float (default: 0.25)
- device:
device for the model (default: 'cuda:0' if available, else 'cpu')
Returns:
- model:
torch model
'''
# Asserting that model is in available models from below
assert base_model in available_base_models
# Accessing the pretrained fair-model
if base_model == 'alexnet':
model = models.alexnet(pretrained=True)
hidden_1_in = 9216
if base_model == 'densenet121':
model = models.densenet121(pretrained=True)
hidden_1_in = 1024
if base_model == 'vgg11':
model = models.vgg11(pretrained=True)
hidden_1_in = 25088
if base_model == 'vgg11_bn':
model = models.vgg11_bn(pretrained=True)
hidden_1_in = 25088
if base_model == 'vgg19':
model = models.vgg19(pretrained=True)
hidden_1_in = 25088
# Freezing the pretrained model
for parameters in model.parameters():
parameters.requires_grad = False
# Customizing the classifier
classifier = nn.Sequential(nn.Linear(hidden_1_in, hidden_1_out), nn.ReLU(),
nn.Dropout(p=dropout_rate),
nn.Linear(hidden_1_out, hidden_2_out),
nn.ReLU(), nn.Dropout(p=dropout_rate),
nn.LogSoftmax(dim=1))
model.classifier = classifier
# Setting the model to device
model.to(device)
return model
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=False) -> (nn.Module, int):
"""get models from https://pytorch.org/hub/.
Args:
model_name (string): model name.
num_classes (int): the output dimension of model classifier.
feature_extract (bool): if true, will freeze all the gradients.
use_pretrained (bool): if true, model will load pretrained weights.
Return:
model, input size.
"""
# Initialize these variables which will be set in this if statement. Each of these variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
elif model_name == "vision_transformer":
""" Vision Transformer base 16
"""
model_ft = models.vit_b_16(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the primary net
num_ftrs = model_ft.hidden_dim
model_ft.heads = nn.Linear(num_ftrs, num_classes)
input_size = 224
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def build(self):
# Transform for input images.
input_size = 299 if self.predictor_name == 'inception_v3' else 224
self.transform = T.Compose([
T.ToPILImage(),
T.Resize((input_size, input_size)),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
if self.predictor_name == 'alexnet':
model = M.alexnet(pretrained=True)
elif self.predictor_name == 'vgg11':
model = M.vgg11(pretrained=True)
elif self.predictor_name == 'vgg13':
model = M.vgg13(pretrained=True)
elif self.predictor_name == 'vgg16':
model = M.vgg16(pretrained=True)
elif self.predictor_name == 'vgg19':
model = M.vgg19(pretrained=True)
elif self.predictor_name == 'vgg11_bn':
model = M.vgg11_bn(pretrained=True)
elif self.predictor_name == 'vgg13_bn':
model = M.vgg13_bn(pretrained=True)
elif self.predictor_name == 'vgg16_bn':
model = M.vgg16_bn(pretrained=True)
elif self.predictor_name == 'vgg19_bn':
model = M.vgg19_bn(pretrained=True)
elif self.predictor_name == 'googlenet':
model = M.googlenet(pretrained=True, aux_logits=False)
elif self.predictor_name == 'inception_v3':
model = M.inception_v3(pretrained=True, aux_logits=False)
elif self.predictor_name == 'resnet18':
model = M.resnet18(pretrained=True)
elif self.predictor_name == 'resnet34':
model = M.resnet34(pretrained=True)
elif self.predictor_name == 'resnet50':
model = M.resnet50(pretrained=True)
elif self.predictor_name == 'resnet101':
model = M.resnet101(pretrained=True)
elif self.predictor_name == 'resnet152':
model = M.resnet152(pretrained=True)
elif self.predictor_name == 'resnext50':
model = M.resnext50_32x4d(pretrained=True)
elif self.predictor_name == 'resnext101':
model = M.resnext101_32x8d(pretrained=True)
elif self.predictor_name == 'wideresnet50':
model = M.wide_resnet50_2(pretrained=True)
elif self.predictor_name == 'wideresnet101':
model = M.wide_resnet101_2(pretrained=True)
elif self.predictor_name == 'densenet121':
model = M.densenet121(pretrained=True)
elif self.predictor_name == 'densenet169':
model = M.densenet169(pretrained=True)
elif self.predictor_name == 'densenet201':
model = M.densenet201(pretrained=True)
elif self.predictor_name == 'densenet161':
model = M.densenet161(pretrained=True)
else:
raise NotImplementedError(f'Unsupported architecture '
f'`{self.predictor_name}`!')
model.eval()
if self.imagenet_logits:
self.net = model
self.feature_dim = (1000, )
return
if self.architecture_type == 'AlexNet':
layers = list(model.features.children())
if not self.spatial_feature:
layers.append(nn.Flatten())
self.feature_dim = (256 * 6 * 6, )
else:
self.feature_dim = (256, 6, 6)
elif self.architecture_type == 'VGG':
layers = list(model.features.children())
if not self.spatial_feature:
layers.append(nn.Flatten())
self.feature_dim = (512 * 7 * 7, )
else:
self.feature_dim = (512, 7, 7)
elif self.architecture_type == 'Inception':
if self.predictor_name == 'googlenet':
final_res = 7
num_channels = 1024
layers = list(model.children())[:-3]
elif self.predictor_name == 'inception_v3':
final_res = 8
num_channels = 2048
layers = list(model.children())[:-1]
layers.insert(3, nn.MaxPool2d(kernel_size=3, stride=2))
layers.insert(6, nn.MaxPool2d(kernel_size=3, stride=2))
else:
raise NotImplementedError(
f'Unsupported Inception architecture '
f'`{self.predictor_name}`!')
if not self.spatial_feature:
layers.append(nn.AdaptiveAvgPool2d((1, 1)))
layers.append(nn.Flatten())
self.feature_dim = (num_channels, )
else:
self.feature_dim = (num_channels, final_res, final_res)
elif self.architecture_type == 'ResNet':
if self.predictor_name in ['resnet18', 'resnet34']:
num_channels = 512
elif self.predictor_name in [
'resnet50', 'resnet101', 'resnet152', 'resnext50',
'resnext101', 'wideresnet50', 'wideresnet101'
]:
num_channels = 2048
else:
raise NotImplementedError(f'Unsupported ResNet architecture '
f'`{self.predictor_name}`!')
if not self.spatial_feature:
layers = list(model.children())[:-1]
layers.append(nn.Flatten())
self.feature_dim = (num_channels, )
else:
layers = list(model.children())[:-2]
self.feature_dim = (num_channels, 7, 7)
elif self.architecture_type == 'DenseNet':
if self.predictor_name == 'densenet121':
num_channels = 1024
elif self.predictor_name == 'densenet169':
num_channels = 1664
elif self.predictor_name == 'densenet201':
num_channels = 1920
elif self.predictor_name == 'densenet161':
num_channels = 2208
else:
raise NotImplementedError(f'Unsupported DenseNet architecture '
f'`{self.predictor_name}`!')
layers = list(model.features.children())
if not self.spatial_feature:
layers.append(nn.ReLU(inplace=True))
layers.append(nn.AdaptiveAvgPool2d((1, 1)))
layers.append(nn.Flatten())
self.feature_dim = (num_channels, )
else:
self.feature_dim = (num_channels, 7, 7)
else:
raise NotImplementedError(f'Unsupported architecture type '
f'`{self.architecture_type}`!')
self.net = nn.Sequential(*layers)
test_dir = 'test'
if not os.path.exists(os.path.join(data_dir, test_dir, 'unknown')):
os.mkdir(os.path.join(data_dir, test_dir, 'unknown'))
for root, dirs, files in os.walk(os.path.join(data_dir, test_dir),
topdown=False):
for file in files:
shutil.move(os.path.join(root, file),
os.path.join(root, 'unknown', file))
pretrained = True
alexnet = models.alexnet(pretrained=pretrained)
vgg11 = models.vgg11(pretrained=pretrained)
vgg11_bn = models.vgg11_bn(pretrained=pretrained)
vgg13 = models.vgg13(pretrained=pretrained)
vgg13_bn = models.vgg13_bn(pretrained=pretrained)
vgg16 = models.vgg16(pretrained=pretrained)
vgg16_bn = models.vgg16_bn(pretrained=pretrained)
vgg19 = models.vgg19(pretrained=pretrained)
vgg19_bn = models.vgg19_bn(pretrained=pretrained)
resnet18 = models.resnet18(pretrained=pretrained)
resnet34 = models.resnet34(pretrained=pretrained)
resnet50 = models.resnet50(pretrained=pretrained)
resnet101 = models.resnet101(pretrained=pretrained)
resnet152 = models.resnet152(pretrained=pretrained)
squeezenet1_0 = models.squeezenet1_0(pretrained=pretrained)
squeezenet1_1 = models.squeezenet1_1(pretrained=pretrained)
def initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
model_ft = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return model_ft, input_size
def initialize_model(model_name, num_classes, feature_extract, verbose=False):
"""Initialize required model and set it up for feature extracting or finetuning.
Args:
model_name: Type of model to initialize.
num_classes: Total number of target classes.
feature_extract: Feature extracting or finetuning.
verbose: Print model info in the end or not.
Returns:
model_ft: Initialized model.
params_to_update: List of parameters to be updated during training.
"""
model_ft = None
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
elif model_name == "squeezenet":
""" Squeezenet
"""
with warnings.catch_warnings(
): # temporarily suppress warnings about deprecated functions
warnings.simplefilter("ignore")
model_ft = models.squeezenet1_0(pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
model_ft.classifier[1] = nn.Conv2d(512,
num_classes,
kernel_size=(1, 1),
stride=(1, 1))
model_ft.num_classes = num_classes
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=True)
set_parameter_requires_grad(model_ft, feature_extract)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
else: # Unreachable
exit()
# Gather the parameters to be optimized
params_to_update = list(
filter(lambda p: p.requires_grad, model_ft.parameters()))
# Print model info
if verbose:
print()
print(model_ft)
print()
print("Params to learn:")
for name, param in model_ft.named_parameters():
if param.requires_grad:
print('\t', name)
return model_ft, params_to_update
num_classes, img_size, po_train_max, _ = dataset_config(args.dataset)
# Adhoc code to make it works on AffMNIST
if args.dataset == 'AffMNIST':
po_train_max = 7
np.random.seed(0)
torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
device = 'cpu'
if __name__ == '__main__':
if args.model == 'vgg11':
model = models.vgg11_bn(num_classes=num_classes).double().to(device)
elif args.model == 'shallowcnn':
model = MNISTNet(3, num_classes, img_size).double().to(device)
if args.model == 'resnet18':
model = models.resnet18(num_classes=num_classes).double().to(device)
torch.save(model.state_dict(), './model_init.pth')
model.eval()
with torch.no_grad():
high.start_counters([
events.PAPI_DP_OPS,
])
x_test_batch = torch.rand(1,
3,
img_size,
img_size,
def _initialize_model(self, model_name, use_pretrained, num_classes, feature_extract, pretrained_model_path):
if model_name == "resnet":
""" Resnet18
"""
model_ft = models.resnet18(pretrained=use_pretrained)
self.set_parameter_requires_grad(feature_extract, model_ft)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "alexnet":
""" Alexnet
"""
model_ft = models.alexnet(pretrained=use_pretrained)
self.set_parameter_requires_grad(feature_extract, model_ft)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
model_ft = models.vgg11_bn(pretrained=use_pretrained)
self.set_parameter_requires_grad(feature_extract, model_ft)
num_ftrs = model_ft.classifier[6].in_features
model_ft.classifier[6] = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
model_ft = models.squeezenet1_0(pretrained=use_pretrained)
self.set_parameter_requires_grad(feature_extract, model_ft)
model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
model_ft.num_classes = num_classes
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
model_ft = models.densenet121(pretrained=use_pretrained)
self.set_parameter_requires_grad(feature_extract, model_ft)
num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
model_ft = models.inception_v3(pretrained=use_pretrained)
self.set_parameter_requires_grad(feature_extract, model_ft)
# Handle the auxilary net
num_ftrs = model_ft.AuxLogits.fc.in_features
model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 299
else:
raise ValueError(f"Invalid model name {model_name}")
if pretrained_model_path is not None:
self._load_pretrained_model(model_ft, pretrained_model_path)
return model_ft, input_size, feature_extract
