def enqueue_blobs(self, gpu_id, blob_names, blobs): """Put a mini-batch on a BlobsQueue.""" #print('------loader.py enqueue_blobs') assert len(blob_names) == len(blobs) t = time.time() dev = c2_utils.CudaDevice(gpu_id) queue_name = 'gpu_{}/{}'.format(gpu_id, self._blobs_queue_name) blob_names = ['gpu_{}/{}'.format(gpu_id, b) for b in blob_names] for (blob_name, blob) in zip(blob_names, blobs): workspace.FeedBlob(blob_name, blob, device_option=dev) logger.debug( 'enqueue_blobs {}: workspace.FeedBlob: {}'. format(gpu_id, time.time() - t) ) t = time.time() op = core.CreateOperator( 'SafeEnqueueBlobs', [queue_name] + blob_names, blob_names + [queue_name + '_enqueue_status'], device_option=dev ) workspace.RunOperatorOnce(op) logger.debug( 'enqueue_blobs {}: workspace.RunOperatorOnce: {}'. format(gpu_id, time.time() - t) )
def setup_model_for_training(model, weights_file, output_dir): """Loaded saved weights and create the network in the C2 workspace.""" logger = logging.getLogger(__name__) add_model_training_inputs(model) if weights_file: # Override random weight initialization with weights from a saved model nu.initialize_gpu_from_weights_file(model, weights_file, gpu_id=0) logger.info("{}".format(cfg.TRAIN.Load_SqueezeNetWeights)) logger.info("{}".format(cfg.TRAIN.SqueezeNetWeightsFile)) # print (cfg.TRAIN.SqueezeNetWeightsFile) if cfg.TRAIN.Load_SqueezeNetWeights: prefix = "gpu_0/" logger.info( '\n\n\n\n========> Loading Weights For SqueezeNet<======================\n\n\n\n' ) pickle_file = cfg.TRAIN.SqueezeNetWeightsFile with open(pickle_file, 'rb') as file: weights = pickle.load(file) dev = c2_utils.CudaDevice(0) for i in weights.keys(): # workspace.FetchBlob(prefix+i) workspace.FeedBlob(prefix + i, weights[i], device_option=dev) # Even if we're randomly initializing we still need to synchronize # parameters across GPUs nu.broadcast_parameters(model) workspace.CreateNet(model.net) logger.info('Outputs saved to: {:s}'.format(os.path.abspath(output_dir))) dump_proto_files(model, output_dir) # Start loading mini-batches and enqueuing blobs model.roi_data_loader.register_sigint_handler() model.roi_data_loader.start(prefill=True) return output_dir