PyTorch version to screen and cleanup (#1325)
* Create flatten_recursive() helper function * cleanup * print torch version
This commit is contained in:
1 parent
81d320109f
commit
19e2482458
4 files changed
+18
-11
No files matched your search
+3
-4
@@ -1,11 +1,10 @@
|
||||
import argparse
|
||||
import logging
|
||||
import math
|
||||
import sys
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
import math
|
||||
|
||||
sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -74,7 +73,7 @@ class Model(nn.Module):
|
||||
|
||||
# Define model
|
||||
if nc and nc != self.yaml['nc']:
|
||||
print('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc))
|
||||
logger.info('Overriding model.yaml nc=%g with nc=%g' % (self.yaml['nc'], nc))
|
||||
self.yaml['nc'] = nc # override yaml value
|
||||
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out
|
||||
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
|
||||
@@ -93,7 +92,7 @@ class Model(nn.Module):
|
||||
# Init weights, biases
|
||||
initialize_weights(self)
|
||||
self.info()
|
||||
print('')
|
||||
logger.info('')
|
||||
|
||||
def forward(self, x, augment=False, profile=False):
|
||||
if augment:
|
||||
|
||||
Reference in new issue
Block a user