tianyi
/
Fast-R-CNN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

shared linear

This commit is contained in:
ClF3 2024-11-19 13:46:49 +08:00
parent b4eb23917e
commit 36651b22de
1 changed files with 12 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
model.py
View File

@ -71,18 +71,16 @@ class FastRCNN(nn.Module):
# hidden_dim -> hidden_dim. #
##############################################################################
# Replace "pass" statement with your code
self.cls_head = nn.Sequential(
self.shared_fc = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.Dropout(drop_ratio),
nn.ReLU(),
nn.Linear(hidden_dim, num_classes+1)
)
self.bbox_head = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.Dropout(drop_ratio),
nn.ReLU(),
nn.Linear(hidden_dim, 4)
nn.Linear(hidden_dim, hidden_dim)
)
self.cls_head = nn.Linear(hidden_dim, self.num_classes+1) # The cls head is a Linear layer that predicts num_classes + 1 (background).
self.bbox_head = nn.Linear(hidden_dim, 4)# The det head is a Linear layer that predicts offsets(dim=4).
##############################################################################
# END OF YOUR CODE #
##############################################################################
@ -139,8 +137,9 @@ class FastRCNN(nn.Module):
# print(feat.shape)
# forward heads, get predicted cls scores & offsets
cls_scores=self.cls_head(feat)
bbox_offsets=self.bbox_head(feat)
shared_feat = self.shared_fc(feat)
cls_scores=self.cls_head(shared_feat)
bbox_offsets=self.bbox_head(shared_feat)
# print(cls_scores.shape, bbox_offsets.shape)
# assign targets with proposals
@ -216,11 +215,11 @@ class FastRCNN(nn.Module):
# perform RoI Pool & mean pool
feat=torchvision.ops.roi_pool(feat, torch.cat((proposal_batch_ids.unsqueeze(1), proposals),dim=1), output_size=(self.roi_output_w, self.roi_output_h))
feat = feat.mean(dim=[2, 3])
shared_feat = self.shared_fc(feat)
# forward heads, get predicted cls scores & offsets
cls_scores = self.cls_head(feat)
cls_scores = self.cls_head(shared_feat)
# print(cls_scores.shape)
bbox_offsets = self.bbox_head(feat)
bbox_offsets = self.bbox_head(shared_feat)
# print(bbox_offsets.shape)
# get predicted boxes & class label & confidence probability
proposals = generate_proposal(proposals, bbox_offsets)