wondervictor commited on
Commit
2b2f933
·
1 Parent(s): 1b32236

add requirements

Browse files
Files changed (1) hide show
  1. condition/midas/midas/vit.py +23 -6
condition/midas/midas/vit.py CHANGED
@@ -383,22 +383,39 @@ def _make_vit_b_rn50_backbone(
383
  pretrained = nn.Module()
384
 
385
  pretrained.model = model
 
386
 
387
  if use_vit_only == True:
388
  pretrained.model.blocks[hooks[0]].register_forward_hook(
389
- get_activation("1"))
390
  pretrained.model.blocks[hooks[1]].register_forward_hook(
391
- get_activation("2"))
392
  else:
393
  pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
394
- get_activation("1"))
395
  pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
396
- get_activation("2"))
397
 
398
  pretrained.model.blocks[hooks[2]].register_forward_hook(
399
- get_activation("3"))
400
  pretrained.model.blocks[hooks[3]].register_forward_hook(
401
- get_activation("4"))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
402
 
403
  pretrained.activations = activations
404
 
 
383
  pretrained = nn.Module()
384
 
385
  pretrained.model = model
386
+ activations = []
387
 
388
  if use_vit_only == True:
389
  pretrained.model.blocks[hooks[0]].register_forward_hook(
390
+ hook_act)
391
  pretrained.model.blocks[hooks[1]].register_forward_hook(
392
+ hook_act)
393
  else:
394
  pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
395
+ hook_act)
396
  pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
397
+ hook_act)
398
 
399
  pretrained.model.blocks[hooks[2]].register_forward_hook(
400
+ hook_act)
401
  pretrained.model.blocks[hooks[3]].register_forward_hook(
402
+ hook_act)
403
+
404
+ # if use_vit_only == True:
405
+ # pretrained.model.blocks[hooks[0]].register_forward_hook(
406
+ # get_activation("1"))
407
+ # pretrained.model.blocks[hooks[1]].register_forward_hook(
408
+ # get_activation("2"))
409
+ # else:
410
+ # pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
411
+ # get_activation("1"))
412
+ # pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
413
+ # get_activation("2"))
414
+
415
+ # pretrained.model.blocks[hooks[2]].register_forward_hook(
416
+ # get_activation("3"))
417
+ # pretrained.model.blocks[hooks[3]].register_forward_hook(
418
+ # get_activation("4"))
419
 
420
  pretrained.activations = activations
421