Skip to content

Commit

Permalink
Added Dropout
Browse files Browse the repository at this point in the history
ultralytics#12 lead to overfitting. We attempt to address this in the added segmentation path
  • Loading branch information
manole-alexandru committed Apr 11, 2023
1 parent ef08e87 commit 00d6490
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 1 deletion.
7 changes: 7 additions & 0 deletions models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -856,16 +856,23 @@ def __init__(self, in_channels):
self.cv2 = Conv(32, 64, k=3)
self.cv3 = Conv(64, 1, act=False)
self.relu = nn.ReLU()

self.dropout_weak= nn.Dropout(0.25)
self.dropout_normal = nn.Dropout(0.5)
# self.sigmoid = nn.Sigmoid()

def forward(self, x):
# print('----entry shape', x.shape, '---\n')
x = self.dropout_weak(x)
x = self.cv1(x)
x = self.upsample(x)
# x = self.relu(x)
# print('----upsample shape', x.shape, '---\n')
x = self.dropout_normal(x)
x = self.cv2(x)
x = self.upsample(x)

x = self.dropout_normal(x)
# x = self.relu(x)
x = self.cv3(x)
# print('----out shape', x.shape, '---\n')
Expand Down
2 changes: 1 addition & 1 deletion train.py
Original file line number Diff line number Diff line change
Expand Up @@ -536,7 +536,7 @@ def parse_opt(known=False):


def main(opt, callbacks=Callbacks()):
print('\n---------- VERSION:', '#0012', '----------\n')
print('\n---------- VERSION:', '#0013', '----------\n')
# Checks
if RANK in {-1, 0}:
print_args(vars(opt))
Expand Down

0 comments on commit 00d6490

Please sign in to comment.