loss = 0
for img in img_all:for i in range(4):for j in range(4):loss_ij = lamda_1*(c_pred-c_label)**2 + c_label*(x_pred-x_label)**2 +c_label*(y_pred-y_label)**2 + c_label*(w_pred-w_label)**2 + c_label*(h_pred-h_label)**2loss += loss_ij
loss.backward()

loss = 0
for img in img_all:for i in range(3):for j in range(4):c_loss = lamda_1*(c_pred-c_label)**2geo_loss = c_label*(x_pred-x_label)**2 +c_label*(y_pred-y_label)**2 + c_label*(w_pred-w_label)**2 + c_label*(h_pred-h_label)**2class_loss = 1/m * mse_loss(p_pred, p_label)loss_ij =c_loss  + geo_loss + class_lossloss += loss_ij
loss.backward()

loss = 0
for img in img_all:for i in range(3):for j in range(4):c_loss = lamda_1*(c_pred-c_label)**2geo_loss = c_label_big*(x_big_pred-x_big_label)**2 +c_label_big*(y_big_pred-y_big_label)**2 + c_label_big*(w_big_pred-w_big_label)**2 + c_label_big*(h_big_pred-h_big_label)**2 +c_label_small*(x_small_pred-x_small_label)**2 +c_label_small*(y_small_pred-y_small_label)**2 + c_label_small*(w_small_pred-w_small_label)**2 + c_label_small*(h_small_pred-h_small_label)**2class_loss = 1/m * mse_loss(p_pred, p_label)loss_ij =c_loss  + geo_loss + class_lossloss += loss_ij
loss.backward()

class VGG(nn.Module):def __init__(self):super(VGG,self).__init__()# the vgg's layers#self.features = featurescfg = [64,64,'M',128,128,'M',256,256,256,'M',512,512,512,'M',512,512,512,'M']layers= []batch_norm = Falsein_channels = 3for v in cfg:if v == 'M':layers += [nn.MaxPool2d(kernel_size=2,stride = 2)]else:conv2d = nn.Conv2d(in_channels,v,kernel_size=3,padding = 1)if batch_norm:layers += [conv2d,nn.Batchnorm2d(v),nn.ReLU(inplace=True)]else:layers += [conv2d,nn.ReLU(inplace=True)]in_channels = v# use the vgg layers to get the featureself.features = nn.Sequential(*layers)# 全局池化self.avgpool = nn.AdaptiveAvgPool2d((7,7))# 决策层：分类层self.classifier = nn.Sequential(nn.Linear(512*7*7,4096),nn.ReLU(True),nn.Dropout(),nn.Linear(4096,4096),nn.ReLU(True),nn.Dropout(),nn.Linear(4096,1000),)for m in self.modules():if isinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,mode='fan_out',nonlinearity='relu')if m.bias is not None: nn.init.constant_(m.bias,0)elif isinstance(m,nn.BatchNorm2d):nn.init.constant_(m.weight,1)nn.init.constant_(m.bias,1)elif isinstance(m,nn.Linear):nn.init.normal_(m.weight,0,0.01)nn.init.constant_(m.bias,0)def forward(self,x):x = self.features(x)x_fea = xx = self.avgpool(x)x_avg = xx = x.view(x.size(0),-1)x = self.classifier(x)return x,x_fea,x_avgdef extractor(self,x):x = self.features(x)return x

       self.detector = nn.Sequential(nn.Linear(512*7*7,4096),nn.ReLU(True),nn.Dropout(),nn.Linear(4096,1470),)

class YOLOV1(nn.Module):def __init__(self):super(YOLOV1,self).__init__()vgg = VGG()self.extractor = vgg.extractorself.avgpool = nn.AdaptiveAvgPool2d((7,7))# 决策层：检测层self.detector = nn.Sequential(nn.Linear(512*7*7,4096),nn.ReLU(True),nn.Dropout(),#nn.Linear(4096,1470),nn.Linear(4096,245),#nn.Linear(4096,5),)for m in self.modules():if isinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,mode='fan_out',nonlinearity='relu')if m.bias is not None: nn.init.constant_(m.bias,0)elif isinstance(m,nn.BatchNorm2d):nn.init.constant_(m.weight,1)nn.init.constant_(m.bias,1)elif isinstance(m,nn.Linear):nn.init.normal_(m.weight,0,0.01)nn.init.constant_(m.bias,0)def forward(self,x):x = self.extractor(x)#import pdb#pdb.set_trace()x = self.avgpool(x)x = x.view(x.size(0),-1)x = self.detector(x)b,_ = x.shape#x = x.view(b,7,7,30)x = x.view(b,7,7,5)#x = x.view(b,1,1,5)return x

if __name__ == '__main__':vgg = VGG()x  = torch.randn(1,3,512,512)feature,x_fea,x_avg = vgg(x)print(feature.shape)print(x_fea.shape)print(x_avg.shape)yolov1 = YOLOV1()feature = yolov1(x)# feature_size b*7*7*30print(feature.shape)

if __name__ == "__main__":train()

def train():for epoch in range(epochs):ts = time.time()for iter, batch in enumerate(train_loader):optimizer.zero_grad()# 取图片inputs = input_process(batch)# 取标注labels = target_process(batch)# 获取得到输出outputs = yolov1_model(inputs)#import pdb#pdb.set_trace()#loss = criterion(outputs, labels)loss,lm,glm,clm = lossfunc_details(outputs,labels)loss.backward()optimizer.step()#print(torch.cat([outputs.detach().view(1,5),labels.view(1,5)],0).view(2,5))if iter % 10 == 0:#    print(torch.cat([outputs.detach().view(1,5),labels.view(1,5)],0).view(2,5))print("epoch{}, iter{}, loss: {}, lr: {}".format(epoch, iter, loss.data.item(),optimizer.state_dict()['param_groups'][0]['lr']))#print("Finish epoch {}, time elapsed {}".format(epoch, time.time() - ts))#print("*"*30)#val(epoch)scheduler.step()

def input_process(batch):#import pdb#pdb.set_trace()batch_size=len(batch[0])input_batch= torch.zeros(batch_size,3,448,448)for i in range(batch_size):inputs_tmp = Variable(batch[0][i])inputs_tmp1=cv2.resize(inputs_tmp.permute([1,2,0]).numpy(),(448,448))inputs_tmp2=torch.tensor(inputs_tmp1).permute([2,0,1])input_batch[i:i+1,:,:,:]= torch.unsqueeze(inputs_tmp2,0)return input_batch 

def target_process(batch,grid_number=7):# batch[1]表示label# batch[0]表示imagebatch_size=len(batch[0])target_batch= torch.zeros(batch_size,grid_number,grid_number,30)#import pdb#pdb.set_trace()for i in range(batch_size):labels = batch[1]batch_labels = labels[i]#import pdb#pdb.set_trace()number_box = len(batch_labels['boxes'])for wi in range(grid_number):for hi in range(grid_number):# 遍历每个标注的框for bi in range(number_box):bbox=batch_labels['boxes'][bi]_,himg,wimg = batch[0][i].numpy().shapebbox = bbox/ torch.tensor([wimg,himg,wimg,himg])#import pdb#pdb.set_trace()center_x= (bbox[0]+bbox[2])*0.5center_y= (bbox[1]+bbox[3])*0.5#print("[%s,%s,%s],[%s,%s,%s]"%(wi/grid_number,center_x,(wi+1)/grid_number,hi/grid_number,center_y,(hi+1)/grid_number))if center_x<=(wi+1)/grid_number and center_x>=wi/grid_number and center_y<=(hi+1)/grid_number and center_y>= hi/grid_number:#pdb.set_trace()cbbox =  torch.cat([torch.ones(1),bbox])# 中心点落在grid内，target_batch[i:i+1,wi:wi+1,hi:hi+1,:] = torch.unsqueeze(cbbox,0)#else:#cbbox =  torch.cat([torch.zeros(1),bbox])#import pdb#pdb.set_trace()#rint(target_batch[i:i+1,wi:wi+1,hi:hi+1,:])#target_batch[i:i+1,wi:wi+1,hi:hi+1,:] = torch.unsqueeze(cbbox,0)return target_batch

def lossfunc_details(outputs,labels):# 判断维度assert ( outputs.shape == labels.shape),"outputs shape[%s] not equal labels shape[%s]"%(outputs.shape,labels.shape)#import pdb#pdb.set_trace()b,w,h,c = outputs.shapeloss = 0#import pdb#pdb.set_trace()conf_loss_matrix = torch.zeros(b,w,h)geo_loss_matrix = torch.zeros(b,w,h)loss_matrix = torch.zeros(b,w,h)for bi in range(b):for wi in range(w):for hi in range(h):#import pdb#pdb.set_trace()# detect_vector=[confidence,x,y,w,h]detect_vector = outputs[bi,wi,hi]gt_dv = labels[bi,wi,hi]conf_pred = detect_vector[0]conf_gt = gt_dv[0]x_pred = detect_vector[1]x_gt = gt_dv[1]y_pred = detect_vector[2]y_gt = gt_dv[2]w_pred = detect_vector[3]w_gt = gt_dv[3]h_pred = detect_vector[4]h_gt = gt_dv[4]loss_confidence = (conf_pred-conf_gt)**2 #loss_geo = (x_pred-x_gt)**2 + (y_pred-y_gt)**2 + (w_pred**0.5-w_gt**0.5)**2 + (h_pred**0.5-h_gt**0.5)**2loss_geo = (x_pred-x_gt)**2 + (y_pred-y_gt)**2 + (w_pred-w_gt)**2 + (h_pred-h_gt)**2loss_geo = conf_gt*loss_geoloss_tmp = loss_confidence + 0.3*loss_geo#print("loss[%s,%s] = %s,%s"%(wi,hi,loss_confidence.item(),loss_geo.item()))loss += loss_tmpconf_loss_matrix[bi,wi,hi]=loss_confidencegeo_loss_matrix[bi,wi,hi]=loss_geoloss_matrix[bi,wi,hi]=loss_tmp#打印出batch中每张片的位置loss,和置信度输出print(geo_loss_matrix)print(outputs[0,:,:,0]>0.5)return loss,loss_matrix,geo_loss_matrix,conf_loss_matrix